From 086803a8ddd59974d4c6680a6fe50372dd606911 Mon Sep 17 00:00:00 2001 From: Yliken <3112951925@qq.com> Date: Fri, 14 Aug 2026 23:41:57 +0800 Subject: [PATCH] first commit --- .dockerignore | 52 + .env.example | 16 + .gitignore | 84 + DEPLOY.md | 306 +++ README.md | 40 + claude code/AGENTS.md | 37 + claude code/Dockerfile | 75 + claude code/RUN.md | 70 + claude code/blackboard-server.mjs | 297 +++ claude code/mcp-servers.json | 9 + claude code/prompts/agent.md | 18 + claude code/supervisor.sh | 215 +++ go.mod | 64 + go.sum | 178 ++ internal/agent/agent.go | 684 +++++++ internal/agent/api_config.go | 258 +++ internal/agent/compress_test.go | 71 + internal/agent/config.go | 137 ++ internal/agent/coop_async.go | 180 ++ internal/agent/coop_test.go | 85 + internal/agent/docker_config.go | 99 + internal/agent/llm.go | 470 +++++ internal/agent/llm_anthropic.go | 298 +++ internal/agent/session.go | 338 ++++ internal/agent/tools.go | 559 ++++++ internal/agent/tools_coop.go | 556 ++++++ internal/agent/tools_docker.go | 320 ++++ internal/agent/tools_test.go | 76 + internal/agent/types.go | 60 + internal/server/auth.go | 80 + internal/server/benchmark.go | 45 + internal/server/handlers_config.go | 111 ++ internal/server/handlers_coop.go | 63 + internal/server/handlers_sessions.go | 89 + internal/server/response.go | 31 + internal/server/router.go | 86 + internal/server/server.go | 19 + internal/server/ws.go | 177 ++ main.go | 53 + pi-coop/AGENTS.md | 29 + pi-coop/Dockerfile | 58 + pi-coop/RUN.md | 127 ++ pi-coop/coop.ts | 262 +++ pi-coop/prompts/agent.md | 23 + pi-coop/supervisor.sh | 207 ++ pigo/.dockerignore | 10 + pigo/.gitignore | 62 + pigo/.goreleaser.yaml | 86 + pigo/LICENSE | 21 + pigo/agent/agent.go | 165 ++ pigo/agent/agent_test.go | 160 ++ pigo/agent/doc.go | 63 + pigo/agent/options.go | 123 ++ pigo/cmd/pigo/main.go | 548 ++++++ pigo/cmd/pigo/main_test.go | 350 ++++ pigo/config.toml.example | 92 + pigo/coop/AGENTS.md | 29 + pigo/coop/Dockerfile | 56 + pigo/coop/README.md | 196 ++ pigo/coop/RUN.md | 107 ++ pigo/coop/build.ps1 | 19 + pigo/coop/prompts/agent.md | 18 + pigo/coop/supervisor.sh | 182 ++ pigo/go.mod | 68 + pigo/go.sum | 173 ++ pigo/install.sh | 117 ++ pigo/internal/agentcore/content.go | 195 ++ pigo/internal/agentcore/event.go | 217 +++ pigo/internal/agentcore/event_stream.go | 121 ++ pigo/internal/agentcore/event_stream_test.go | 121 ++ pigo/internal/agentcore/helpers.go | 60 + pigo/internal/agentcore/helpers_test.go | 141 ++ pigo/internal/agentcore/hooks.go | 44 + pigo/internal/agentcore/message.go | 210 ++ pigo/internal/agentcore/progress_ctx.go | 22 + pigo/internal/agentcore/progress_ctx_test.go | 51 + pigo/internal/agentcore/tool.go | 61 + pigo/internal/agentcore/types_test.go | 182 ++ pigo/internal/agenttool/bash_background.go | 167 ++ .../agenttool/bash_background_test.go | 134 ++ pigo/internal/agenttool/bash_control.go | 158 ++ pigo/internal/agenttool/bash_tool.go | 328 ++++ pigo/internal/agenttool/bash_tool_test.go | 272 +++ pigo/internal/agenttool/batch_executor.go | 86 + .../internal/agenttool/batch_executor_test.go | 227 +++ pigo/internal/agenttool/blackboard_tool.go | 340 ++++ .../agenttool/blackboard_tool_test.go | 200 ++ pigo/internal/agenttool/edit_tool.go | 232 +++ pigo/internal/agenttool/edit_tool_test.go | 177 ++ pigo/internal/agenttool/file_snapshot.go | 213 +++ pigo/internal/agenttool/file_snapshot_test.go | 152 ++ pigo/internal/agenttool/goal_tool.go | 374 ++++ pigo/internal/agenttool/goal_tool_test.go | 188 ++ pigo/internal/agenttool/htmlmarkdown.go | 48 + pigo/internal/agenttool/htmlmarkdown_test.go | 71 + pigo/internal/agenttool/memory_tool.go | 170 ++ pigo/internal/agenttool/memory_tool_test.go | 179 ++ pigo/internal/agenttool/read_tool.go | 178 ++ pigo/internal/agenttool/read_tool_test.go | 204 ++ pigo/internal/agenttool/registry.go | 220 +++ pigo/internal/agenttool/registry_test.go | 138 ++ pigo/internal/agenttool/search_tool.go | 492 +++++ pigo/internal/agenttool/search_tool_test.go | 225 +++ pigo/internal/agenttool/todo_tool.go | 186 ++ pigo/internal/agenttool/todo_tool_test.go | 143 ++ pigo/internal/agenttool/tool_executor.go | 341 ++++ pigo/internal/agenttool/tool_executor_test.go | 465 +++++ pigo/internal/agenttool/tool_retry.go | 124 ++ pigo/internal/agenttool/webfetch_tool.go | 231 +++ pigo/internal/agenttool/webfetch_tool_test.go | 197 ++ pigo/internal/agenttool/websearch_backends.go | 275 +++ pigo/internal/agenttool/websearch_tool.go | 217 +++ .../internal/agenttool/websearch_tool_test.go | 160 ++ pigo/internal/agenttool/write_tool.go | 125 ++ pigo/internal/agenttool/write_tool_test.go | 164 ++ pigo/internal/builtinskills/bootstrap.go | 177 ++ pigo/internal/builtinskills/bootstrap_test.go | 164 ++ pigo/internal/builtinskills/manifest.go | 74 + .../skills/architecture-diagram/SKILL.md | 163 ++ .../architecture-diagram/assets/template.html | 319 ++++ .../skills/code-to-spec/SKILL.md | 341 ++++ .../builtinskills/skills/graph/SKILL.md | 329 ++++ .../skills/graph/scripts/render_graph_html.py | 179 ++ .../skills/insight-diagram/SKILL.md | 273 +++ .../references/extraction-strategy.md | 62 + .../insight-diagram/scripts/review_svg.py | 432 +++++ .../builtinskills/skills/loop-it/SKILL.md | 562 ++++++ .../builtinskills/skills/modern-go/SKILL.md | 1139 +++++++++++ .../builtinskills/skills/note-it/SKILL.md | 188 ++ .../builtinskills/skills/prd-to-spec/SKILL.md | 409 ++++ .../internal/builtinskills/skills/prd/LICENSE | 21 + .../builtinskills/skills/prd/README.md | 43 + .../builtinskills/skills/prd/SKILL.md | 278 +++ .../skills/prd/test-prompts.json | 27 + .../builtinskills/skills/refactor/SKILL.md | 674 +++++++ .../builtinskills/skills/review-it/SKILL.md | 163 ++ .../skills/review-it/scripts/review-it | 248 +++ .../builtinskills/skills/ship-it/SKILL.md | 181 ++ .../builtinskills/skills/smell/README.md | 39 + .../builtinskills/skills/smell/SKILL.md | 690 +++++++ .../skills/smell/test-prompts.json | 50 + .../builtinskills/skills/to-design/SKILL.md | 281 +++ .../builtinskills/skills/to-issues/SKILL.md | 229 +++ .../builtinskills/skills/weather/SKILL.md | 49 + .../builtinskills/skills/weather/_meta.json | 6 + pigo/internal/cli/btw/btw.go | 302 +++ pigo/internal/cli/btw/btw_config.go | 140 ++ pigo/internal/cli/btw/btw_config_test.go | 146 ++ pigo/internal/cli/config/config.go | 128 ++ pigo/internal/cli/config/config_test.go | 176 ++ pigo/internal/cli/config/memory.go | 265 +++ pigo/internal/cli/config/memory_test.go | 234 +++ pigo/internal/cli/doc.go | 23 + pigo/internal/cli/goal/goal.go | 425 +++++ pigo/internal/cli/goal/goal_test.go | 141 ++ pigo/internal/cli/headless/headless.go | 205 ++ pigo/internal/cli/headless/headless_test.go | 41 + pigo/internal/cli/headless/session.go | 185 ++ pigo/internal/cli/headless/session_test.go | 156 ++ pigo/internal/cli/headless/subagent_rpc.go | 161 ++ .../cli/headless/subagent_rpc_test.go | 120 ++ pigo/internal/cli/host.go | 91 + pigo/internal/cli/liveconfig.go | 38 + pigo/internal/cli/memstatus/memstatus.go | 195 ++ pigo/internal/cli/memstatus/memstatus_test.go | 53 + pigo/internal/cli/persist.go | 53 + pigo/internal/cli/pkgcmd/pkgcmd.go | 152 ++ pigo/internal/cli/prompts/presets_test.go | 33 + pigo/internal/cli/prompts/prompts_cli_test.go | 116 ++ .../cli/prompts/prompts_config_test.go | 115 ++ pigo/internal/cli/prompts/prompts_dir_test.go | 91 + .../cli/prompts/prompts_project_test.go | 139 ++ pigo/internal/cli/prompts/registry.go | 387 ++++ pigo/internal/cli/prompts/think_test.go | 87 + pigo/internal/cli/providerhelp.go | 31 + pigo/internal/cli/providerhelp_test.go | 32 + .../cli/repl/autocomplete_label_test.go | 62 + pigo/internal/cli/repl/btw_help_test.go | 58 + pigo/internal/cli/repl/btw_isolation_test.go | 148 ++ pigo/internal/cli/repl/btw_test.go | 171 ++ pigo/internal/cli/repl/color_test.go | 41 + pigo/internal/cli/repl/dream_repl.go | 240 +++ pigo/internal/cli/repl/dream_repl_test.go | 255 +++ pigo/internal/cli/repl/dream_startup.go | 58 + pigo/internal/cli/repl/dream_startup_test.go | 132 ++ pigo/internal/cli/repl/help_line_test.go | 61 + pigo/internal/cli/repl/host.go | 53 + pigo/internal/cli/repl/interactive.go | 261 +++ pigo/internal/cli/repl/line_editor.go | 766 ++++++++ pigo/internal/cli/repl/line_editor_test.go | 599 ++++++ .../internal/cli/repl/plugin_commands_test.go | 254 +++ pigo/internal/cli/repl/remotecontrol.go | 237 +++ pigo/internal/cli/repl/repl.go | 1150 +++++++++++ pigo/internal/cli/repl/repl_test.go | 656 +++++++ pigo/internal/cli/repl/rewind.go | 160 ++ pigo/internal/cli/repl/rewind_test.go | 89 + pigo/internal/cli/repl/skills_test.go | 165 ++ pigo/internal/cli/repl/status_repl_test.go | 127 ++ pigo/internal/cli/run/hooks_converge_test.go | 106 ++ pigo/internal/cli/run/hooks_driver.go | 64 + pigo/internal/cli/run/hooks_install.go | 223 +++ pigo/internal/cli/run/hooks_install_test.go | 108 ++ pigo/internal/cli/run/hooks_prompt.go | 55 + pigo/internal/cli/run/hooks_prompt_test.go | 97 + pigo/internal/cli/run/hooks_session.go | 51 + pigo/internal/cli/run/hooks_session_test.go | 76 + pigo/internal/cli/run/hooks_stop.go | 105 + pigo/internal/cli/run/hooks_stop_test.go | 84 + pigo/internal/cli/run/hooks_tooluse_test.go | 189 ++ pigo/internal/cli/run/memory_wiring_test.go | 87 + pigo/internal/cli/run/prompt_flags_test.go | 88 + pigo/internal/cli/run/run.go | 604 ++++++ pigo/internal/cli/run/task_wiring_test.go | 39 + pigo/internal/cli/run/thinking_test.go | 120 ++ pigo/internal/cli/run/toolpolicy.go | 214 +++ .../internal/cli/run/toolpolicy_setup_test.go | 212 +++ pigo/internal/cli/run/toolpolicy_test.go | 176 ++ pigo/internal/cli/status/fakehost_test.go | 49 + pigo/internal/cli/status/status.go | 292 +++ pigo/internal/cli/status/status_e2e_test.go | 139 ++ pigo/internal/cli/status/status_test.go | 254 +++ pigo/internal/cli/telemetry.go | 135 ++ pigo/internal/cli/telemetry_test.go | 209 ++ pigo/internal/cli/testutil/prompts.go | 24 + pigo/internal/cli/tui/banner.go | 105 + pigo/internal/cli/tui/banner_test.go | 90 + pigo/internal/cli/tui/bridge.go | 139 ++ pigo/internal/cli/tui/bridge_test.go | 181 ++ pigo/internal/cli/tui/clipimage.go | 122 ++ pigo/internal/cli/tui/doc.go | 14 + pigo/internal/cli/tui/gitinfo.go | 136 ++ pigo/internal/cli/tui/gitinfo_test.go | 99 + pigo/internal/cli/tui/host.go | 91 + pigo/internal/cli/tui/input.go | 167 ++ pigo/internal/cli/tui/input_test.go | 205 ++ pigo/internal/cli/tui/markdown.go | 122 ++ pigo/internal/cli/tui/model.go | 1360 +++++++++++++ pigo/internal/cli/tui/model_test.go | 519 +++++ pigo/internal/cli/tui/msgs.go | 82 + pigo/internal/cli/tui/options.go | 59 + pigo/internal/cli/tui/remotecontrol.go | 206 ++ pigo/internal/cli/tui/remotecontrol_test.go | 145 ++ pigo/internal/cli/tui/run.go | 24 + pigo/internal/cli/tui/selection.go | 112 ++ pigo/internal/cli/tui/selection_test.go | 59 + pigo/internal/cli/tui/session.go | 491 +++++ pigo/internal/cli/tui/session_test.go | 210 ++ pigo/internal/cli/tui/slash.go | 206 ++ pigo/internal/cli/tui/slash_test.go | 216 +++ pigo/internal/cli/tui/spinner.go | 185 ++ pigo/internal/cli/tui/spinner_test.go | 123 ++ pigo/internal/cli/tui/status_test.go | 219 +++ pigo/internal/cli/tui/statusbar.go | 397 ++++ pigo/internal/cli/tui/statusbar_test.go | 164 ++ pigo/internal/cli/tui/subagentpanel.go | 317 ++++ pigo/internal/cli/tui/subagentpanel_test.go | 292 +++ pigo/internal/cli/tui/theme.go | 195 ++ pigo/internal/cli/tui/theme_test.go | 116 ++ pigo/internal/cli/tui/toolcard.go | 195 ++ pigo/internal/cli/tui/toolcard_test.go | 164 ++ pigo/internal/cli/tui/transcript.go | 404 ++++ pigo/internal/cli/tui/transcript_test.go | 324 ++++ pigo/internal/cli/ui/color.go | 52 + pigo/internal/cli/ui/color_test.go | 27 + pigo/internal/cli/ui/imageref.go | 119 ++ pigo/internal/cli/ui/imageref_test.go | 122 ++ pigo/internal/cli/ui/markdown.go | 72 + pigo/internal/cli/ui/markdown_test.go | 23 + pigo/internal/cli/ui/toolrender.go | 58 + pigo/internal/cli/ui/width.go | 16 + pigo/internal/cli/ui/width_test.go | 33 + pigo/internal/clipboard/clipboard.go | 74 + pigo/internal/clipboard/clipboard_test.go | 61 + pigo/internal/compaction/compact.go | 133 ++ pigo/internal/compaction/cutpoint.go | 97 + pigo/internal/compaction/cutpoint_test.go | 119 ++ pigo/internal/compaction/summary.go | 367 ++++ pigo/internal/compaction/summary_test.go | 274 +++ pigo/internal/compaction/tokens.go | 189 ++ pigo/internal/compaction/tokens_test.go | 189 ++ pigo/internal/dream/apply.go | 247 +++ pigo/internal/dream/apply_test.go | 143 ++ pigo/internal/dream/config.go | 48 + pigo/internal/dream/config_test.go | 61 + pigo/internal/dream/consolidator.go | 434 +++++ pigo/internal/dream/consolidator_test.go | 195 ++ pigo/internal/dream/distill.go | 411 ++++ pigo/internal/dream/distill_test.go | 290 +++ pigo/internal/dream/lock.go | 147 ++ pigo/internal/dream/lock_test.go | 191 ++ pigo/internal/dream/plan.go | 359 ++++ pigo/internal/dream/plan_test.go | 216 +++ pigo/internal/dream/prompt.go | 90 + .../dream/reconcile_validation_test.go | 105 + pigo/internal/dream/report.go | 28 + pigo/internal/dream/report_test.go | 66 + pigo/internal/dream/runner.go | 527 +++++ pigo/internal/dream/runner_test.go | 303 +++ pigo/internal/dream/scheduler.go | 103 + pigo/internal/dream/scheduler_test.go | 211 ++ pigo/internal/dream/state.go | 99 + pigo/internal/dream/state_test.go | 100 + pigo/internal/hooks/config.go | 83 + pigo/internal/hooks/config_test.go | 47 + pigo/internal/hooks/dispatch.go | 91 + pigo/internal/hooks/dispatch_test.go | 112 ++ pigo/internal/hooks/matcher.go | 85 + pigo/internal/hooks/matcher_test.go | 111 ++ pigo/internal/hooks/notifier.go | 118 ++ pigo/internal/hooks/notifier_test.go | 122 ++ pigo/internal/hooks/protocol.go | 82 + pigo/internal/hooks/protocol_test.go | 80 + pigo/internal/hooks/runner.go | 162 ++ pigo/internal/hooks/runner_test.go | 130 ++ pigo/internal/jsonrpc/message.go | 111 ++ pigo/internal/jsonrpc/transport.go | 293 +++ pigo/internal/jsonrpc/transport_test.go | 223 +++ pigo/internal/memory/count_test.go | 57 + pigo/internal/memory/ftsquery.go | 62 + pigo/internal/memory/ftsquery_test.go | 55 + pigo/internal/memory/paths.go | 233 +++ pigo/internal/memory/paths_test.go | 203 ++ pigo/internal/memory/reconcile.go | 235 +++ pigo/internal/memory/reconcile_test.go | 256 +++ pigo/internal/memory/schema.go | 66 + pigo/internal/memory/search.go | 173 ++ pigo/internal/memory/search_test.go | 152 ++ pigo/internal/memory/store.go | 137 ++ pigo/internal/memory/store_test.go | 175 ++ pigo/internal/pihost/embed.go | 17 + pigo/internal/pihost/host_e2e_test.go | 156 ++ pigo/internal/pihost/pihost.mjs | 641 +++++++ pigo/internal/pkgmgr/classify.go | 150 ++ pigo/internal/pkgmgr/classify_test.go | 205 ++ pigo/internal/pkgmgr/distribute.go | 260 +++ pigo/internal/pkgmgr/distribute_prompt.go | 85 + .../internal/pkgmgr/distribute_prompt_test.go | 121 ++ pigo/internal/pkgmgr/distribute_skill.go | 48 + pigo/internal/pkgmgr/distribute_skill_test.go | 76 + pigo/internal/pkgmgr/distribute_test.go | 264 +++ pigo/internal/pkgmgr/distribute_theme.go | 38 + pigo/internal/pkgmgr/distribute_theme_test.go | 56 + pigo/internal/pkgmgr/fetch.go | 218 +++ pigo/internal/pkgmgr/fetch_test.go | 165 ++ pigo/internal/pkgmgr/install.go | 122 ++ pigo/internal/pkgmgr/install_test.go | 140 ++ pigo/internal/pkgmgr/layout.go | 110 ++ pigo/internal/pkgmgr/layout_test.go | 58 + pigo/internal/pkgmgr/lockfile.go | 183 ++ pigo/internal/pkgmgr/lockfile_test.go | 141 ++ pigo/internal/pkgmgr/ref.go | 148 ++ pigo/internal/pkgmgr/ref_test.go | 126 ++ pigo/internal/pkgmgr/uninstall.go | 72 + pigo/internal/pkgmgr/uninstall_test.go | 142 ++ pigo/internal/pkgmgr/update.go | 86 + pigo/internal/pkgmgr/update_test.go | 141 ++ pigo/internal/plugin/events.go | 92 + pigo/internal/plugin/events_test.go | 262 +++ pigo/internal/plugin/manager.go | 140 ++ pigo/internal/plugin/manager_test.go | 137 ++ pigo/internal/plugin/manifest.go | 111 ++ pigo/internal/plugin/manifest_test.go | 81 + pigo/internal/plugin/plugin.go | 196 ++ pigo/internal/plugin/plugin_test.go | 320 ++++ pigo/internal/provider/anthropic.go | 334 ++++ pigo/internal/provider/anthropic_test.go | 299 +++ pigo/internal/provider/auth.go | 249 +++ pigo/internal/provider/auth_test.go | 255 +++ pigo/internal/provider/image_test.go | 157 ++ pigo/internal/provider/infer.go | 78 + pigo/internal/provider/infer_test.go | 96 + pigo/internal/provider/openai.go | 245 +++ pigo/internal/provider/openai_test.go | 190 ++ pigo/internal/provider/presets.go | 200 ++ .../internal/provider/presets_catalog_test.go | 118 ++ pigo/internal/provider/presets_test.go | 66 + pigo/internal/provider/protocol.go | 77 + pigo/internal/provider/protocol_test.go | 55 + pigo/internal/provider/provider.go | 118 ++ pigo/internal/provider/provider_interface.go | 79 + .../provider/provider_interface_test.go | 87 + pigo/internal/provider/providers.go | 775 ++++++++ .../provider/providers_anthropic_test.go | 275 +++ .../provider/providers_openai_test.go | 151 ++ pigo/internal/provider/providers_test.go | 227 +++ pigo/internal/provider/registry.go | 377 ++++ pigo/internal/provider/registry_test.go | 164 ++ pigo/internal/provider/resolve.go | 223 +++ pigo/internal/provider/resolve_test.go | 296 +++ pigo/internal/provider/responses.go | 474 +++++ pigo/internal/provider/responses_test.go | 756 ++++++++ pigo/internal/provider/special_auth.go | 225 +++ pigo/internal/provider/special_auth_test.go | 291 +++ pigo/internal/provider/thinking_test.go | 264 +++ pigo/internal/provider/transport.go | 393 ++++ pigo/internal/provider/transport_test.go | 251 +++ pigo/internal/remotecontrol/bridge.go | 180 ++ pigo/internal/remotecontrol/bridge_test.go | 192 ++ pigo/internal/remotecontrol/lanaddr.go | 116 ++ pigo/internal/remotecontrol/lanaddr_test.go | 182 ++ pigo/internal/remotecontrol/protocol.go | 47 + pigo/internal/remotecontrol/qr.go | 51 + pigo/internal/remotecontrol/qr_test.go | 60 + pigo/internal/remotecontrol/server.go | 602 ++++++ pigo/internal/remotecontrol/server_test.go | 464 +++++ pigo/internal/remotecontrol/spa_test.go | 53 + pigo/internal/remotecontrol/token.go | 130 ++ pigo/internal/remotecontrol/token_test.go | 121 ++ pigo/internal/remotecontrol/web/app.js | 153 ++ pigo/internal/remotecontrol/web/index.html | 135 ++ pigo/internal/runtime/args.go | 34 + pigo/internal/runtime/args_test.go | 69 + pigo/internal/runtime/checkpoint.go | 193 ++ pigo/internal/runtime/checkpoint_test.go | 138 ++ pigo/internal/runtime/compaction_test.go | 206 ++ pigo/internal/runtime/config.go | 186 ++ pigo/internal/runtime/config_test.go | 240 +++ pigo/internal/runtime/e2e_robustness_test.go | 149 ++ pigo/internal/runtime/faux_provider_test.go | 429 +++++ pigo/internal/runtime/headless.go | 259 +++ pigo/internal/runtime/headless_test.go | 292 +++ pigo/internal/runtime/loop.go | 503 +++++ pigo/internal/runtime/loop_onstop_test.go | 63 + pigo/internal/runtime/loop_test.go | 286 +++ pigo/internal/runtime/memory_reminder.go | 185 ++ pigo/internal/runtime/memory_reminder_test.go | 150 ++ pigo/internal/runtime/orchestration_test.go | 638 +++++++ pigo/internal/runtime/progress_test.go | 149 ++ pigo/internal/runtime/prompt.go | 203 ++ pigo/internal/runtime/prompt_test.go | 315 +++ pigo/internal/runtime/rebuild.go | 165 ++ pigo/internal/runtime/rebuild_test.go | 194 ++ pigo/internal/runtime/reminder.go | 264 +++ pigo/internal/runtime/reminder_test.go | 194 ++ pigo/internal/runtime/render.go | 86 + pigo/internal/runtime/render_test.go | 164 ++ pigo/internal/runtime/skills.go | 374 ++++ pigo/internal/runtime/slashcommand.go | 538 ++++++ pigo/internal/runtime/slashcommand_test.go | 182 ++ pigo/internal/runtime/stream_response.go | 183 ++ pigo/internal/runtime/stream_response_test.go | 176 ++ pigo/internal/runtime/subagent.go | 472 +++++ .../internal/runtime/subagent_process_test.go | 328 ++++ pigo/internal/runtime/task.go | 141 ++ pigo/internal/runtime/task_test.go | 233 +++ pigo/internal/runtime/telemetry.go | 137 ++ pigo/internal/runtime/telemetry_test.go | 276 +++ pigo/internal/runtime/template.go | 182 ++ pigo/internal/runtime/template_test.go | 170 ++ pigo/internal/runtime/testtools_test.go | 52 + pigo/internal/selfupdate/cache.go | 102 + pigo/internal/selfupdate/cache_test.go | 52 + pigo/internal/selfupdate/update.go | 264 +++ pigo/internal/selfupdate/update_test.go | 175 ++ pigo/internal/selfupdate/version.go | 142 ++ pigo/internal/selfupdate/version_test.go | 142 ++ pigo/internal/session/export.go | 112 ++ pigo/internal/session/export_html.go | 125 ++ pigo/internal/session/export_test.go | 197 ++ pigo/internal/session/inherit.go | 56 + pigo/internal/session/session.go | 700 +++++++ pigo/internal/session/session_test.go | 843 ++++++++ pigo/internal/trust/interactive.go | 305 +++ pigo/internal/trust/interactive_test.go | 319 ++++ pigo/internal/trust/manager.go | 303 +++ pigo/internal/trust/manager_test.go | 338 ++++ tencent/Dockerfile | 202 ++ tencent/build.sh | 58 + web/static/app.js | 1673 ++++++++++++++++ web/static/index.html | 219 +++ web/static/style.css | 1689 +++++++++++++++++ 471 files changed, 91938 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 DEPLOY.md create mode 100644 README.md create mode 100644 claude code/AGENTS.md create mode 100644 claude code/Dockerfile create mode 100644 claude code/RUN.md create mode 100644 claude code/blackboard-server.mjs create mode 100644 claude code/mcp-servers.json create mode 100644 claude code/prompts/agent.md create mode 100644 claude code/supervisor.sh create mode 100644 go.mod create mode 100644 go.sum create mode 100644 internal/agent/agent.go create mode 100644 internal/agent/api_config.go create mode 100644 internal/agent/compress_test.go create mode 100644 internal/agent/config.go create mode 100644 internal/agent/coop_async.go create mode 100644 internal/agent/coop_test.go create mode 100644 internal/agent/docker_config.go create mode 100644 internal/agent/llm.go create mode 100644 internal/agent/llm_anthropic.go create mode 100644 internal/agent/session.go create mode 100644 internal/agent/tools.go create mode 100644 internal/agent/tools_coop.go create mode 100644 internal/agent/tools_docker.go create mode 100644 internal/agent/tools_test.go create mode 100644 internal/agent/types.go create mode 100644 internal/server/auth.go create mode 100644 internal/server/benchmark.go create mode 100644 internal/server/handlers_config.go create mode 100644 internal/server/handlers_coop.go create mode 100644 internal/server/handlers_sessions.go create mode 100644 internal/server/response.go create mode 100644 internal/server/router.go create mode 100644 internal/server/server.go create mode 100644 internal/server/ws.go create mode 100644 main.go create mode 100644 pi-coop/AGENTS.md create mode 100644 pi-coop/Dockerfile create mode 100644 pi-coop/RUN.md create mode 100644 pi-coop/coop.ts create mode 100644 pi-coop/prompts/agent.md create mode 100644 pi-coop/supervisor.sh create mode 100644 pigo/.dockerignore create mode 100644 pigo/.gitignore create mode 100644 pigo/.goreleaser.yaml create mode 100644 pigo/LICENSE create mode 100644 pigo/agent/agent.go create mode 100644 pigo/agent/agent_test.go create mode 100644 pigo/agent/doc.go create mode 100644 pigo/agent/options.go create mode 100644 pigo/cmd/pigo/main.go create mode 100644 pigo/cmd/pigo/main_test.go create mode 100644 pigo/config.toml.example create mode 100644 pigo/coop/AGENTS.md create mode 100644 pigo/coop/Dockerfile create mode 100644 pigo/coop/README.md create mode 100644 pigo/coop/RUN.md create mode 100644 pigo/coop/build.ps1 create mode 100644 pigo/coop/prompts/agent.md create mode 100644 pigo/coop/supervisor.sh create mode 100644 pigo/go.mod create mode 100644 pigo/go.sum create mode 100644 pigo/install.sh create mode 100644 pigo/internal/agentcore/content.go create mode 100644 pigo/internal/agentcore/event.go create mode 100644 pigo/internal/agentcore/event_stream.go create mode 100644 pigo/internal/agentcore/event_stream_test.go create mode 100644 pigo/internal/agentcore/helpers.go create mode 100644 pigo/internal/agentcore/helpers_test.go create mode 100644 pigo/internal/agentcore/hooks.go create mode 100644 pigo/internal/agentcore/message.go create mode 100644 pigo/internal/agentcore/progress_ctx.go create mode 100644 pigo/internal/agentcore/progress_ctx_test.go create mode 100644 pigo/internal/agentcore/tool.go create mode 100644 pigo/internal/agentcore/types_test.go create mode 100644 pigo/internal/agenttool/bash_background.go create mode 100644 pigo/internal/agenttool/bash_background_test.go create mode 100644 pigo/internal/agenttool/bash_control.go create mode 100644 pigo/internal/agenttool/bash_tool.go create mode 100644 pigo/internal/agenttool/bash_tool_test.go create mode 100644 pigo/internal/agenttool/batch_executor.go create mode 100644 pigo/internal/agenttool/batch_executor_test.go create mode 100644 pigo/internal/agenttool/blackboard_tool.go create mode 100644 pigo/internal/agenttool/blackboard_tool_test.go create mode 100644 pigo/internal/agenttool/edit_tool.go create mode 100644 pigo/internal/agenttool/edit_tool_test.go create mode 100644 pigo/internal/agenttool/file_snapshot.go create mode 100644 pigo/internal/agenttool/file_snapshot_test.go create mode 100644 pigo/internal/agenttool/goal_tool.go create mode 100644 pigo/internal/agenttool/goal_tool_test.go create mode 100644 pigo/internal/agenttool/htmlmarkdown.go create mode 100644 pigo/internal/agenttool/htmlmarkdown_test.go create mode 100644 pigo/internal/agenttool/memory_tool.go create mode 100644 pigo/internal/agenttool/memory_tool_test.go create mode 100644 pigo/internal/agenttool/read_tool.go create mode 100644 pigo/internal/agenttool/read_tool_test.go create mode 100644 pigo/internal/agenttool/registry.go create mode 100644 pigo/internal/agenttool/registry_test.go create mode 100644 pigo/internal/agenttool/search_tool.go create mode 100644 pigo/internal/agenttool/search_tool_test.go create mode 100644 pigo/internal/agenttool/todo_tool.go create mode 100644 pigo/internal/agenttool/todo_tool_test.go create mode 100644 pigo/internal/agenttool/tool_executor.go create mode 100644 pigo/internal/agenttool/tool_executor_test.go create mode 100644 pigo/internal/agenttool/tool_retry.go create mode 100644 pigo/internal/agenttool/webfetch_tool.go create mode 100644 pigo/internal/agenttool/webfetch_tool_test.go create mode 100644 pigo/internal/agenttool/websearch_backends.go create mode 100644 pigo/internal/agenttool/websearch_tool.go create mode 100644 pigo/internal/agenttool/websearch_tool_test.go create mode 100644 pigo/internal/agenttool/write_tool.go create mode 100644 pigo/internal/agenttool/write_tool_test.go create mode 100644 pigo/internal/builtinskills/bootstrap.go create mode 100644 pigo/internal/builtinskills/bootstrap_test.go create mode 100644 pigo/internal/builtinskills/manifest.go create mode 100644 pigo/internal/builtinskills/skills/architecture-diagram/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/architecture-diagram/assets/template.html create mode 100644 pigo/internal/builtinskills/skills/code-to-spec/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/graph/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/graph/scripts/render_graph_html.py create mode 100644 pigo/internal/builtinskills/skills/insight-diagram/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/insight-diagram/references/extraction-strategy.md create mode 100644 pigo/internal/builtinskills/skills/insight-diagram/scripts/review_svg.py create mode 100644 pigo/internal/builtinskills/skills/loop-it/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/modern-go/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/note-it/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/prd-to-spec/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/prd/LICENSE create mode 100644 pigo/internal/builtinskills/skills/prd/README.md create mode 100644 pigo/internal/builtinskills/skills/prd/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/prd/test-prompts.json create mode 100644 pigo/internal/builtinskills/skills/refactor/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/review-it/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/review-it/scripts/review-it create mode 100644 pigo/internal/builtinskills/skills/ship-it/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/smell/README.md create mode 100644 pigo/internal/builtinskills/skills/smell/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/smell/test-prompts.json create mode 100644 pigo/internal/builtinskills/skills/to-design/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/to-issues/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/weather/SKILL.md create mode 100644 pigo/internal/builtinskills/skills/weather/_meta.json create mode 100644 pigo/internal/cli/btw/btw.go create mode 100644 pigo/internal/cli/btw/btw_config.go create mode 100644 pigo/internal/cli/btw/btw_config_test.go create mode 100644 pigo/internal/cli/config/config.go create mode 100644 pigo/internal/cli/config/config_test.go create mode 100644 pigo/internal/cli/config/memory.go create mode 100644 pigo/internal/cli/config/memory_test.go create mode 100644 pigo/internal/cli/doc.go create mode 100644 pigo/internal/cli/goal/goal.go create mode 100644 pigo/internal/cli/goal/goal_test.go create mode 100644 pigo/internal/cli/headless/headless.go create mode 100644 pigo/internal/cli/headless/headless_test.go create mode 100644 pigo/internal/cli/headless/session.go create mode 100644 pigo/internal/cli/headless/session_test.go create mode 100644 pigo/internal/cli/headless/subagent_rpc.go create mode 100644 pigo/internal/cli/headless/subagent_rpc_test.go create mode 100644 pigo/internal/cli/host.go create mode 100644 pigo/internal/cli/liveconfig.go create mode 100644 pigo/internal/cli/memstatus/memstatus.go create mode 100644 pigo/internal/cli/memstatus/memstatus_test.go create mode 100644 pigo/internal/cli/persist.go create mode 100644 pigo/internal/cli/pkgcmd/pkgcmd.go create mode 100644 pigo/internal/cli/prompts/presets_test.go create mode 100644 pigo/internal/cli/prompts/prompts_cli_test.go create mode 100644 pigo/internal/cli/prompts/prompts_config_test.go create mode 100644 pigo/internal/cli/prompts/prompts_dir_test.go create mode 100644 pigo/internal/cli/prompts/prompts_project_test.go create mode 100644 pigo/internal/cli/prompts/registry.go create mode 100644 pigo/internal/cli/prompts/think_test.go create mode 100644 pigo/internal/cli/providerhelp.go create mode 100644 pigo/internal/cli/providerhelp_test.go create mode 100644 pigo/internal/cli/repl/autocomplete_label_test.go create mode 100644 pigo/internal/cli/repl/btw_help_test.go create mode 100644 pigo/internal/cli/repl/btw_isolation_test.go create mode 100644 pigo/internal/cli/repl/btw_test.go create mode 100644 pigo/internal/cli/repl/color_test.go create mode 100644 pigo/internal/cli/repl/dream_repl.go create mode 100644 pigo/internal/cli/repl/dream_repl_test.go create mode 100644 pigo/internal/cli/repl/dream_startup.go create mode 100644 pigo/internal/cli/repl/dream_startup_test.go create mode 100644 pigo/internal/cli/repl/help_line_test.go create mode 100644 pigo/internal/cli/repl/host.go create mode 100644 pigo/internal/cli/repl/interactive.go create mode 100644 pigo/internal/cli/repl/line_editor.go create mode 100644 pigo/internal/cli/repl/line_editor_test.go create mode 100644 pigo/internal/cli/repl/plugin_commands_test.go create mode 100644 pigo/internal/cli/repl/remotecontrol.go create mode 100644 pigo/internal/cli/repl/repl.go create mode 100644 pigo/internal/cli/repl/repl_test.go create mode 100644 pigo/internal/cli/repl/rewind.go create mode 100644 pigo/internal/cli/repl/rewind_test.go create mode 100644 pigo/internal/cli/repl/skills_test.go create mode 100644 pigo/internal/cli/repl/status_repl_test.go create mode 100644 pigo/internal/cli/run/hooks_converge_test.go create mode 100644 pigo/internal/cli/run/hooks_driver.go create mode 100644 pigo/internal/cli/run/hooks_install.go create mode 100644 pigo/internal/cli/run/hooks_install_test.go create mode 100644 pigo/internal/cli/run/hooks_prompt.go create mode 100644 pigo/internal/cli/run/hooks_prompt_test.go create mode 100644 pigo/internal/cli/run/hooks_session.go create mode 100644 pigo/internal/cli/run/hooks_session_test.go create mode 100644 pigo/internal/cli/run/hooks_stop.go create mode 100644 pigo/internal/cli/run/hooks_stop_test.go create mode 100644 pigo/internal/cli/run/hooks_tooluse_test.go create mode 100644 pigo/internal/cli/run/memory_wiring_test.go create mode 100644 pigo/internal/cli/run/prompt_flags_test.go create mode 100644 pigo/internal/cli/run/run.go create mode 100644 pigo/internal/cli/run/task_wiring_test.go create mode 100644 pigo/internal/cli/run/thinking_test.go create mode 100644 pigo/internal/cli/run/toolpolicy.go create mode 100644 pigo/internal/cli/run/toolpolicy_setup_test.go create mode 100644 pigo/internal/cli/run/toolpolicy_test.go create mode 100644 pigo/internal/cli/status/fakehost_test.go create mode 100644 pigo/internal/cli/status/status.go create mode 100644 pigo/internal/cli/status/status_e2e_test.go create mode 100644 pigo/internal/cli/status/status_test.go create mode 100644 pigo/internal/cli/telemetry.go create mode 100644 pigo/internal/cli/telemetry_test.go create mode 100644 pigo/internal/cli/testutil/prompts.go create mode 100644 pigo/internal/cli/tui/banner.go create mode 100644 pigo/internal/cli/tui/banner_test.go create mode 100644 pigo/internal/cli/tui/bridge.go create mode 100644 pigo/internal/cli/tui/bridge_test.go create mode 100644 pigo/internal/cli/tui/clipimage.go create mode 100644 pigo/internal/cli/tui/doc.go create mode 100644 pigo/internal/cli/tui/gitinfo.go create mode 100644 pigo/internal/cli/tui/gitinfo_test.go create mode 100644 pigo/internal/cli/tui/host.go create mode 100644 pigo/internal/cli/tui/input.go create mode 100644 pigo/internal/cli/tui/input_test.go create mode 100644 pigo/internal/cli/tui/markdown.go create mode 100644 pigo/internal/cli/tui/model.go create mode 100644 pigo/internal/cli/tui/model_test.go create mode 100644 pigo/internal/cli/tui/msgs.go create mode 100644 pigo/internal/cli/tui/options.go create mode 100644 pigo/internal/cli/tui/remotecontrol.go create mode 100644 pigo/internal/cli/tui/remotecontrol_test.go create mode 100644 pigo/internal/cli/tui/run.go create mode 100644 pigo/internal/cli/tui/selection.go create mode 100644 pigo/internal/cli/tui/selection_test.go create mode 100644 pigo/internal/cli/tui/session.go create mode 100644 pigo/internal/cli/tui/session_test.go create mode 100644 pigo/internal/cli/tui/slash.go create mode 100644 pigo/internal/cli/tui/slash_test.go create mode 100644 pigo/internal/cli/tui/spinner.go create mode 100644 pigo/internal/cli/tui/spinner_test.go create mode 100644 pigo/internal/cli/tui/status_test.go create mode 100644 pigo/internal/cli/tui/statusbar.go create mode 100644 pigo/internal/cli/tui/statusbar_test.go create mode 100644 pigo/internal/cli/tui/subagentpanel.go create mode 100644 pigo/internal/cli/tui/subagentpanel_test.go create mode 100644 pigo/internal/cli/tui/theme.go create mode 100644 pigo/internal/cli/tui/theme_test.go create mode 100644 pigo/internal/cli/tui/toolcard.go create mode 100644 pigo/internal/cli/tui/toolcard_test.go create mode 100644 pigo/internal/cli/tui/transcript.go create mode 100644 pigo/internal/cli/tui/transcript_test.go create mode 100644 pigo/internal/cli/ui/color.go create mode 100644 pigo/internal/cli/ui/color_test.go create mode 100644 pigo/internal/cli/ui/imageref.go create mode 100644 pigo/internal/cli/ui/imageref_test.go create mode 100644 pigo/internal/cli/ui/markdown.go create mode 100644 pigo/internal/cli/ui/markdown_test.go create mode 100644 pigo/internal/cli/ui/toolrender.go create mode 100644 pigo/internal/cli/ui/width.go create mode 100644 pigo/internal/cli/ui/width_test.go create mode 100644 pigo/internal/clipboard/clipboard.go create mode 100644 pigo/internal/clipboard/clipboard_test.go create mode 100644 pigo/internal/compaction/compact.go create mode 100644 pigo/internal/compaction/cutpoint.go create mode 100644 pigo/internal/compaction/cutpoint_test.go create mode 100644 pigo/internal/compaction/summary.go create mode 100644 pigo/internal/compaction/summary_test.go create mode 100644 pigo/internal/compaction/tokens.go create mode 100644 pigo/internal/compaction/tokens_test.go create mode 100644 pigo/internal/dream/apply.go create mode 100644 pigo/internal/dream/apply_test.go create mode 100644 pigo/internal/dream/config.go create mode 100644 pigo/internal/dream/config_test.go create mode 100644 pigo/internal/dream/consolidator.go create mode 100644 pigo/internal/dream/consolidator_test.go create mode 100644 pigo/internal/dream/distill.go create mode 100644 pigo/internal/dream/distill_test.go create mode 100644 pigo/internal/dream/lock.go create mode 100644 pigo/internal/dream/lock_test.go create mode 100644 pigo/internal/dream/plan.go create mode 100644 pigo/internal/dream/plan_test.go create mode 100644 pigo/internal/dream/prompt.go create mode 100644 pigo/internal/dream/reconcile_validation_test.go create mode 100644 pigo/internal/dream/report.go create mode 100644 pigo/internal/dream/report_test.go create mode 100644 pigo/internal/dream/runner.go create mode 100644 pigo/internal/dream/runner_test.go create mode 100644 pigo/internal/dream/scheduler.go create mode 100644 pigo/internal/dream/scheduler_test.go create mode 100644 pigo/internal/dream/state.go create mode 100644 pigo/internal/dream/state_test.go create mode 100644 pigo/internal/hooks/config.go create mode 100644 pigo/internal/hooks/config_test.go create mode 100644 pigo/internal/hooks/dispatch.go create mode 100644 pigo/internal/hooks/dispatch_test.go create mode 100644 pigo/internal/hooks/matcher.go create mode 100644 pigo/internal/hooks/matcher_test.go create mode 100644 pigo/internal/hooks/notifier.go create mode 100644 pigo/internal/hooks/notifier_test.go create mode 100644 pigo/internal/hooks/protocol.go create mode 100644 pigo/internal/hooks/protocol_test.go create mode 100644 pigo/internal/hooks/runner.go create mode 100644 pigo/internal/hooks/runner_test.go create mode 100644 pigo/internal/jsonrpc/message.go create mode 100644 pigo/internal/jsonrpc/transport.go create mode 100644 pigo/internal/jsonrpc/transport_test.go create mode 100644 pigo/internal/memory/count_test.go create mode 100644 pigo/internal/memory/ftsquery.go create mode 100644 pigo/internal/memory/ftsquery_test.go create mode 100644 pigo/internal/memory/paths.go create mode 100644 pigo/internal/memory/paths_test.go create mode 100644 pigo/internal/memory/reconcile.go create mode 100644 pigo/internal/memory/reconcile_test.go create mode 100644 pigo/internal/memory/schema.go create mode 100644 pigo/internal/memory/search.go create mode 100644 pigo/internal/memory/search_test.go create mode 100644 pigo/internal/memory/store.go create mode 100644 pigo/internal/memory/store_test.go create mode 100644 pigo/internal/pihost/embed.go create mode 100644 pigo/internal/pihost/host_e2e_test.go create mode 100644 pigo/internal/pihost/pihost.mjs create mode 100644 pigo/internal/pkgmgr/classify.go create mode 100644 pigo/internal/pkgmgr/classify_test.go create mode 100644 pigo/internal/pkgmgr/distribute.go create mode 100644 pigo/internal/pkgmgr/distribute_prompt.go create mode 100644 pigo/internal/pkgmgr/distribute_prompt_test.go create mode 100644 pigo/internal/pkgmgr/distribute_skill.go create mode 100644 pigo/internal/pkgmgr/distribute_skill_test.go create mode 100644 pigo/internal/pkgmgr/distribute_test.go create mode 100644 pigo/internal/pkgmgr/distribute_theme.go create mode 100644 pigo/internal/pkgmgr/distribute_theme_test.go create mode 100644 pigo/internal/pkgmgr/fetch.go create mode 100644 pigo/internal/pkgmgr/fetch_test.go create mode 100644 pigo/internal/pkgmgr/install.go create mode 100644 pigo/internal/pkgmgr/install_test.go create mode 100644 pigo/internal/pkgmgr/layout.go create mode 100644 pigo/internal/pkgmgr/layout_test.go create mode 100644 pigo/internal/pkgmgr/lockfile.go create mode 100644 pigo/internal/pkgmgr/lockfile_test.go create mode 100644 pigo/internal/pkgmgr/ref.go create mode 100644 pigo/internal/pkgmgr/ref_test.go create mode 100644 pigo/internal/pkgmgr/uninstall.go create mode 100644 pigo/internal/pkgmgr/uninstall_test.go create mode 100644 pigo/internal/pkgmgr/update.go create mode 100644 pigo/internal/pkgmgr/update_test.go create mode 100644 pigo/internal/plugin/events.go create mode 100644 pigo/internal/plugin/events_test.go create mode 100644 pigo/internal/plugin/manager.go create mode 100644 pigo/internal/plugin/manager_test.go create mode 100644 pigo/internal/plugin/manifest.go create mode 100644 pigo/internal/plugin/manifest_test.go create mode 100644 pigo/internal/plugin/plugin.go create mode 100644 pigo/internal/plugin/plugin_test.go create mode 100644 pigo/internal/provider/anthropic.go create mode 100644 pigo/internal/provider/anthropic_test.go create mode 100644 pigo/internal/provider/auth.go create mode 100644 pigo/internal/provider/auth_test.go create mode 100644 pigo/internal/provider/image_test.go create mode 100644 pigo/internal/provider/infer.go create mode 100644 pigo/internal/provider/infer_test.go create mode 100644 pigo/internal/provider/openai.go create mode 100644 pigo/internal/provider/openai_test.go create mode 100644 pigo/internal/provider/presets.go create mode 100644 pigo/internal/provider/presets_catalog_test.go create mode 100644 pigo/internal/provider/presets_test.go create mode 100644 pigo/internal/provider/protocol.go create mode 100644 pigo/internal/provider/protocol_test.go create mode 100644 pigo/internal/provider/provider.go create mode 100644 pigo/internal/provider/provider_interface.go create mode 100644 pigo/internal/provider/provider_interface_test.go create mode 100644 pigo/internal/provider/providers.go create mode 100644 pigo/internal/provider/providers_anthropic_test.go create mode 100644 pigo/internal/provider/providers_openai_test.go create mode 100644 pigo/internal/provider/providers_test.go create mode 100644 pigo/internal/provider/registry.go create mode 100644 pigo/internal/provider/registry_test.go create mode 100644 pigo/internal/provider/resolve.go create mode 100644 pigo/internal/provider/resolve_test.go create mode 100644 pigo/internal/provider/responses.go create mode 100644 pigo/internal/provider/responses_test.go create mode 100644 pigo/internal/provider/special_auth.go create mode 100644 pigo/internal/provider/special_auth_test.go create mode 100644 pigo/internal/provider/thinking_test.go create mode 100644 pigo/internal/provider/transport.go create mode 100644 pigo/internal/provider/transport_test.go create mode 100644 pigo/internal/remotecontrol/bridge.go create mode 100644 pigo/internal/remotecontrol/bridge_test.go create mode 100644 pigo/internal/remotecontrol/lanaddr.go create mode 100644 pigo/internal/remotecontrol/lanaddr_test.go create mode 100644 pigo/internal/remotecontrol/protocol.go create mode 100644 pigo/internal/remotecontrol/qr.go create mode 100644 pigo/internal/remotecontrol/qr_test.go create mode 100644 pigo/internal/remotecontrol/server.go create mode 100644 pigo/internal/remotecontrol/server_test.go create mode 100644 pigo/internal/remotecontrol/spa_test.go create mode 100644 pigo/internal/remotecontrol/token.go create mode 100644 pigo/internal/remotecontrol/token_test.go create mode 100644 pigo/internal/remotecontrol/web/app.js create mode 100644 pigo/internal/remotecontrol/web/index.html create mode 100644 pigo/internal/runtime/args.go create mode 100644 pigo/internal/runtime/args_test.go create mode 100644 pigo/internal/runtime/checkpoint.go create mode 100644 pigo/internal/runtime/checkpoint_test.go create mode 100644 pigo/internal/runtime/compaction_test.go create mode 100644 pigo/internal/runtime/config.go create mode 100644 pigo/internal/runtime/config_test.go create mode 100644 pigo/internal/runtime/e2e_robustness_test.go create mode 100644 pigo/internal/runtime/faux_provider_test.go create mode 100644 pigo/internal/runtime/headless.go create mode 100644 pigo/internal/runtime/headless_test.go create mode 100644 pigo/internal/runtime/loop.go create mode 100644 pigo/internal/runtime/loop_onstop_test.go create mode 100644 pigo/internal/runtime/loop_test.go create mode 100644 pigo/internal/runtime/memory_reminder.go create mode 100644 pigo/internal/runtime/memory_reminder_test.go create mode 100644 pigo/internal/runtime/orchestration_test.go create mode 100644 pigo/internal/runtime/progress_test.go create mode 100644 pigo/internal/runtime/prompt.go create mode 100644 pigo/internal/runtime/prompt_test.go create mode 100644 pigo/internal/runtime/rebuild.go create mode 100644 pigo/internal/runtime/rebuild_test.go create mode 100644 pigo/internal/runtime/reminder.go create mode 100644 pigo/internal/runtime/reminder_test.go create mode 100644 pigo/internal/runtime/render.go create mode 100644 pigo/internal/runtime/render_test.go create mode 100644 pigo/internal/runtime/skills.go create mode 100644 pigo/internal/runtime/slashcommand.go create mode 100644 pigo/internal/runtime/slashcommand_test.go create mode 100644 pigo/internal/runtime/stream_response.go create mode 100644 pigo/internal/runtime/stream_response_test.go create mode 100644 pigo/internal/runtime/subagent.go create mode 100644 pigo/internal/runtime/subagent_process_test.go create mode 100644 pigo/internal/runtime/task.go create mode 100644 pigo/internal/runtime/task_test.go create mode 100644 pigo/internal/runtime/telemetry.go create mode 100644 pigo/internal/runtime/telemetry_test.go create mode 100644 pigo/internal/runtime/template.go create mode 100644 pigo/internal/runtime/template_test.go create mode 100644 pigo/internal/runtime/testtools_test.go create mode 100644 pigo/internal/selfupdate/cache.go create mode 100644 pigo/internal/selfupdate/cache_test.go create mode 100644 pigo/internal/selfupdate/update.go create mode 100644 pigo/internal/selfupdate/update_test.go create mode 100644 pigo/internal/selfupdate/version.go create mode 100644 pigo/internal/selfupdate/version_test.go create mode 100644 pigo/internal/session/export.go create mode 100644 pigo/internal/session/export_html.go create mode 100644 pigo/internal/session/export_test.go create mode 100644 pigo/internal/session/inherit.go create mode 100644 pigo/internal/session/session.go create mode 100644 pigo/internal/session/session_test.go create mode 100644 pigo/internal/trust/interactive.go create mode 100644 pigo/internal/trust/interactive_test.go create mode 100644 pigo/internal/trust/manager.go create mode 100644 pigo/internal/trust/manager_test.go create mode 100644 tencent/Dockerfile create mode 100644 tencent/build.sh create mode 100644 web/static/app.js create mode 100644 web/static/index.html create mode 100644 web/static/style.css diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..6a4c8a2 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,52 @@ +# .dockerignore — 构建上下文只保留镜像必需的文件 +# 构建上下文 = 项目根目录,Dockerfile = tencent/Dockerfile + +# ---- 排除整个目录 ---- +.git/ +.gocache/ +.idea/ +.workbuddy/ +claude code/ +coop_blackboard/ +data/ +demo/ +pi/ +pigo/ +internal/ + +# ---- 排除根目录 Go 源码与编译产物 ---- +*.go +go.mod +go.sum +main.go +agent.exe +agent-linux-amd64 +nul + +# ---- 排除 tencent 下的非必需文件 ---- +tencent/agent.exe +tencent/agent.tar.gz* +tencent/log.txt + +# ---- 排除脚本与文档 ---- +*.py +*.md +*.txt +deploy_*.py +log_analyzer.py +test_log_analyzer.py +sort_script.py +response.txt +pigo-coop.tar + +# ---- 排除配置与密钥 ---- +.env +.env.example + +# ---- 排除杂项 ---- +*.tar +*.tar.gz +*.log +.tmp_* +__pycache__/ +node_modules/ diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ef11cb8 --- /dev/null +++ b/.env.example @@ -0,0 +1,16 @@ +SILICONFLOW_API_KEY=your_api_key_here +SILICONFLOW_BASE_URL=https://api.siliconflow.cn/v1/chat/completions +AGENT_MODEL=Qwen/Qwen3-32B +AGENT_PORT=8080 +# 访问授权码(默认空=关闭)。Docker 云端部署时设置该值即启用: +# 访问者必须在 Web 界面输入正确授权码才能使用(REST API 与 WebSocket 同样受保护) +AGENT_AUTH_CODE= +AGENT_MAX_CONTEXT_TOKENS=28000 +AGENT_KEEP_RECENT_MESSAGES=12 +AGENT_MAX_TOOL_RESULT_CHARS=12000 +AGENT_MAX_ITERATIONS=12 +AGENT_TOOL_TIMEOUT_SECONDS=120 +AGENT_REQUEST_TIMEOUT_SECONDS=180 +AGENT_STREAM_OUTPUT_TOKENS=4096 +AGENT_COMPACTION_TOKENS=1024 +AGENT_TEMPERATURE=0.3 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7290ee1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,84 @@ +# ============================================================ +# 敏感配置(严禁入库) +# ============================================================ +# 环境变量文件(含 API Key / 授权码) +.env +.env.* +!.env.example + +# ============================================================ +# 运行时数据(本地运行 / run_coop 产生) +# ============================================================ +# 会话、api-config.json、docker-config.json、sessions.json +data/ +# 协作黑板目录(run_coop / pi-coop 本地测试产生) +blackboard/ +coop_blackboard/ + +# ============================================================ +# Go 编译产物(控制层 / pigo) +# ============================================================ +# go build . 在仓库根产生的二进制(Linux/macOS 无扩展名); +# 锚定 /agent,避免误伤 internal/agent/ 与 pigo/agent/ 源码目录 +/agent +agent.exe +agent-linux-* +# pigo 交叉编译输出(pigo/coop/tmp/pigo-linux-amd64) +pigo/coop/tmp/ +# Go 缓存目录 +.gocache/ +.gomodcache/ + +# ============================================================ +# Docker 镜像导出与压缩包 +# ============================================================ +# *.tar 不匹配 agent.tar.gz,两个都要有 +*.tar +*.tar.gz + +# ============================================================ +# 安全工具二进制(体积大,不入库;见 DEPLOY.md 需自行准备) +# ============================================================ +tencent/tools/ + +# ============================================================ +# 日志与临时文件 +# ============================================================ +*.log +log.txt +.tmp_* +*.tmp +nul + +# ============================================================ +# Python +# ============================================================ +__pycache__/ +*.py[cod] +.venv/ +venv/ + +# ============================================================ +# Node(pi-coop 本地调试可能产生) +# ============================================================ +node_modules/ +dist/ + +# ============================================================ +# IDE / 编辑器 +# ============================================================ +.idea/ +.vscode/ +*.swp + +# ============================================================ +# 操作系统 +# ============================================================ +.DS_Store +Thumbs.db +Desktop.ini + +# ============================================================ +# AI 工具记忆目录 +# ============================================================ +.workbuddy/ diff --git a/DEPLOY.md b/DEPLOY.md new file mode 100644 index 0000000..e6f2542 --- /dev/null +++ b/DEPLOY.md @@ -0,0 +1,306 @@ +# 部署文档 + +本文档覆盖 Local Agent 的全部部署方式:本地运行、Docker 协作引擎镜像、Tsecbench 托管模式镜像,以及完整的配置项说明与常见问题排障。 + +--- + +## 1. 架构概览 + +``` +┌────────────────────────────────────────────────┐ +│ 控制层 agent(Go 二进制,本项目根目录) │ +│ - HTTP/WS 服务(Gin),serve web/static 前端 │ +│ - 多轮工具调用循环(LLM 驱动) │ +│ - run_coop 工具:调起协作容器异步执行子任务 │ +└──────────────┬─────────────────────────────────┘ + │ docker run / 子进程 + ┌───────────┼─────────────┬───────────────┐ + ▼ ▼ ▼ ▼ + pi-coop pigo-coop claude-coop supervisor.sh +(pi 引擎) (pigo 引擎) (Claude Code) (本地子进程模式, + 默认 Tsecbench 托管用) +``` + +| 组件 | 目录 | 说明 | +| --- | --- | --- | +| 控制层 | `main.go` + `internal/` | Agent 核心、HTTP/WS 服务、工具实现 | +| 前端 | `web/static/` | 由控制层直接 serve,无需单独部署 | +| pi 引擎 | `pi-coop/` | **默认协作引擎**,基于 npm 包 `@earendil-works/pi-coding-agent` | +| pigo 引擎 | `pigo/` | Go 自研引擎,需交叉编译二进制后构建镜像 | +| claude 引擎 | `claude code/` | 包装 Claude Code CLI,仅支持 anthropic 协议 | +| 托管镜像 | `tencent/` | Tsecbench 平台一体化镜像(控制层 + worker 本地子进程模式) | + +--- + +## 2. 环境要求 + +| 依赖 | 版本 | 用途 | +| --- | --- | --- | +| Go | ≥ 1.22(以 go.mod 为准) | 编译控制层 / pigo | +| Docker | 任意现代版本 | 构建协作镜像 / run_coop 工具 | +| Node | ≥ 22.19 | pi 引擎运行时(仅容器内需要,宿主机不用装) | +| WSL(Windows) | — | 在 Windows 上构建 Linux 镜像时使用 | +| Python 3 + paramiko | 可选 | 仅 `deploy_*.py` 远程部署脚本需要 | + +--- + +## 3. 配置说明 + +### 3.1 配置优先级 + +``` +环境变量 > .env 文件 > data/api-config.json > 内置默认值 +``` + +- `.env` 放在程序工作目录,启动时自动加载(已存在的环境变量优先)。 +- `api-config.json` 由 Web 设置页写入,存于数据目录(`AGENT_DATA_DIR`)。 +- 容器部署时通过 `docker run -e` 注入的环境变量具有最高优先级。 + +### 3.2 LLM 配置(协作引擎共用) + +| 环境变量 | 必填 | 默认 | 说明 | +| --- | --- | --- | --- | +| `LLM_API_KEY` | ✅ | — | LLM API 密钥 | +| `LLM_BASE_URL` | ✅ | — | 接口地址。anthropic 协议会自动剥离 `/v1`、`/v1/messages`、`/messages` 后缀 | +| `LLM_MODEL` | ✅ | — | 模型名,需与供应商支持的格式完全一致(如 `deepseek-v4-pro` / `deepseek-v4-flash`) | +| `LLM_PROVIDER` | ✅ | `openai` | `openai` / `anthropic` | +| `LLM_ENGINE` | 否 | `pi` | 协作引擎:`pi` / `pigo` / `claude` | +| `COOP_MODE` | 否 | `local` | 协作模式(`local` 为本地子进程模式) | + +DeepSeek Anthropic 端点示例: + +``` +LLM_BASE_URL=https://api.deepseek.com/anthropic +LLM_PROVIDER=anthropic +LLM_MODEL=deepseek-v4-flash +``` + +### 3.3 控制层运行配置 + +| 环境变量 | 内置默认 | 说明 | +| --- | --- | --- | +| `AGENT_PORT` | `8080` | Web 服务端口 | +| `AGENT_WORKSPACE` | 当前目录 | Agent 可读写的根目录 | +| `AGENT_DATA_DIR` | `/data` | 会话与配置存储目录 | +| `AGENT_AUTH_CODE` | 空(关闭) | 访问授权码,云端部署强烈建议设置 | +| `AGENT_MAX_CONTEXT_TOKENS` | `1000000` | 超过该估算值触发历史压缩 | +| `AGENT_KEEP_RECENT_MESSAGES` | `50` | 压缩时保留的最近消息数 | +| `AGENT_MAX_TOOL_RESULT_CHARS` | `100000` | 单次工具结果最大字符数 | +| `AGENT_MAX_ITERATIONS` | `50` | 单轮最多工具调用次数 | +| `AGENT_TOOL_TIMEOUT_SECONDS` | `600` | 单次工具执行超时 | +| `AGENT_REQUEST_TIMEOUT_SECONDS` | `600` | 单次 LLM 请求超时 | +| `AGENT_STREAM_OUTPUT_TOKENS` | `65536` | 流式输出 max_tokens(防截断关键项) | +| `AGENT_COMPACTION_TOKENS` | `16384` | 压缩摘要最大输出 token | +| `AGENT_TEMPERATURE` | `0.3` | 采样温度 | +| `AGENT_DOCKER_SOCKET` | 平台默认 | Docker 连接地址(本地 unix socket / npipe / 远程) | + +> 注:`.env.example` 中的数值是保守示例。复杂任务(长思考、大输出)请参照上表默认值,尤其是 `AGENT_STREAM_OUTPUT_TOKENS`——过低会导致 LLM 思考阶段耗尽配额、工具调用被截断。 + +### 3.4 托管模式开关 + +| 环境变量 | 说明 | +| --- | --- | +| `BENCHMARK_TOKEN` | 非空时进入 Tsecbench 托管模式:启动即自动开跑评测,无需外部触发;为空则是普通本地/云部署 | + +### 3.5 遗留变量 + +`SILICONFLOW_API_KEY` / `SILICONFLOW_BASE_URL` / `AGENT_MODEL` 为早期硅基流动专用配置,仍向后兼容;新部署统一使用 `LLM_*` 系列。 + +--- + +## 4. 部署方式一:本地直接运行 + +### 4.1 开发运行 + +```powershell +go run . +``` + +打开 http://localhost:8080 ,在 Web 设置页填入 API Key(写入 `data/api-config.json`),或提前配置 `.env`。 + +### 4.2 生产运行(Linux + systemd) + +1. 交叉编译(Windows 上): + +```powershell +$env:GOOS="linux"; $env:GOARCH="amd64"; $env:CGO_ENABLED="0" +go build -trimpath -ldflags="-s -w" -o agent-linux-amd64 . +``` + +2. 上传 `agent-linux-amd64` 与 `web/static/` 到服务器 `/opt/agent/`: + +``` +/opt/agent/agent # 二进制(chmod 755) +/opt/agent/web/static/ # 前端 +/opt/agent/.env # 配置(可选,也可全部用环境变量) +``` + +3. systemd 单元 `/etc/systemd/system/agent-web.service`: + +```ini +[Unit] +Description=Local Agent Web +After=network.target + +[Service] +WorkingDirectory=/opt/agent +Environment=LLM_API_KEY=sk-xxxx +Environment=LLM_BASE_URL=https://api.deepseek.com/anthropic +Environment=LLM_MODEL=deepseek-v4-flash +Environment=LLM_PROVIDER=anthropic +ExecStart=/opt/agent/agent +Restart=always + +[Install] +WantedBy=multi-user.target +``` + +```bash +systemctl daemon-reload && systemctl enable --now agent-web +curl -s http://localhost:8080/api/v1/health # 验证 +``` + +--- + +## 5. 部署方式二:Docker 协作引擎 + +控制层通过 `run_coop` 工具调起协作容器。**镜像名与容器名前缀是代码内置约定,不可更改**: + +| 引擎 | 镜像名 | 容器名前缀 | +| --- | --- | --- | +| pi(默认) | `pi-coop` | `pi-coop-` | +| pigo | `pigo-coop` | `pigo-coop-` | +| claude | `claude-coop` | `claude-coop-` | + +### 5.1 构建 pi-coop(默认引擎,推荐) + +```bash +# 仓库根目录,构建上下文为 pi-coop/,无需预编译任何二进制 +docker build -f pi-coop/Dockerfile -t pi-coop pi-coop/ +``` + +镜像内含 node 22 + npm 全局安装的 pi coding agent + 系统/Python 工具。 +运行参数(由 run_coop 自动注入,也可手动测试): + +```bash +docker run --rm \ + -e MODEL=deepseek-v4-flash \ + -e BASE_URL=https://api.deepseek.com/anthropic \ + -e API_KEY= \ + -e PROTOCOL=anthropic \ + -e TASK="任务描述" \ + -e ROUND_MAX=6 \ + pi-coop +``` + +### 5.2 构建 pigo-coop + +需先交叉编译 pigo 二进制(构建上下文为 `pigo/`): + +```bash +cd pigo +GOOS=linux CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" \ + -o coop/tmp/pigo-linux-amd64 ./cmd/pigo +docker build -f coop/Dockerfile -t pigo-coop . +``` + +### 5.3 构建 claude-coop + +```bash +docker build -f "claude code/Dockerfile" -t claude-coop "claude code/" +``` + +仅支持 anthropic 协议;openai 协议请求会被拒绝,请切换引擎或更换端点。 + +### 5.4 切换引擎 + +Web 设置页修改协作引擎,或设置 `LLM_ENGINE=pigo` / `LLM_ENGINE=claude`。切换前确保对应镜像已构建。 + +--- + +## 6. 部署方式三:Tsecbench 托管模式(tencent/) + +一体化镜像 = 控制层 agent + pi worker(本地子进程模式)+ 安全工具集,适合上传到 Tsecbench 评测平台。 + +### 6.1 前置准备 + +镜像构建依赖以下文件,**均不在 git 仓库内,需自行准备**: + +| 文件 | 说明 | +| --- | --- | +| `tencent/agent-linux-amd64` | 控制层交叉编译产物(见 4.2 步骤 1) | +| `tencent/tools/nuclei` | 漏洞扫描器静态二进制 | +| `tencent/tools/observer_ward` | 指纹识别静态二进制 | +| `tencent/tools/chisel` | 内网隧道静态二进制 | +| `tencent/tools/nuclei-templates/` | Nuclei 模板库 | +| `tencent/tools/FingerprintHub-defaultv4/plugins/` | observer_ward 指纹规则库 | + +### 6.2 构建与导出 + +在 **WSL** 中执行(脚本内 PROJECT_ROOT 按实际路径调整): + +```bash +bash tencent/build.sh +# 或自定义镜像名:IMAGE_NAME=myagent bash tencent/build.sh +``` + +脚本流程:检查二进制 → `docker build`(上下文为项目根,`.dockerignore` 排除无关文件)→ `docker save | gzip` 导出。 + +产物:`tencent/agent.tar.gz`(约 350MB),直接上传 Tsecbench 平台即可。 + +### 6.3 本地验证镜像 + +```bash +docker run --rm -p 8081:8080 \ + -e LLM_API_KEY=sk-xxxx \ + -e LLM_BASE_URL=https://api.deepseek.com/anthropic \ + -e LLM_MODEL=deepseek-v4-flash \ + -e LLM_PROVIDER=anthropic \ + tsecbench-agent:latest +``` + +- 打开 http://localhost:8081 验证 Web 界面。 +- 平台托管时由运行时注入 `BENCHMARK_TOKEN`,agent 检测到后自动启动评测流程。 +- 镜像内不含 `data/` 目录与任何预置配置,全部由环境变量注入。 + +--- + +## 7. 验证与排障 + +### 7.1 健康检查 + +```bash +curl http://localhost:8080/api/v1/health # 服务存活(始终放行,不受授权码保护) +curl http://localhost:8080/api/v1/llm/config # 查看生效的 LLM 配置(需授权码) +``` + +### 7.2 常见问题 + +| 现象 | 原因与解决 | +| --- | --- | +| 侧栏显示「模型未配置」 | 未设置 `LLM_API_KEY`,也未在 Web 设置页填写。配置后重启或刷新 | +| LLM 请求 404 | Base URL 拼接问题。anthropic 协议不要手动加 `/v1`,程序会自动处理后缀;openai 协议填到根地址即可 | +| 工具调用 JSON 被截断、任务失败 | `AGENT_STREAM_OUTPUT_TOKENS` 过低,LLM 思考阶段耗尽 max_tokens。保持默认 65536 | +| run_coop 报找不到镜像 | 对应引擎镜像未构建(见第 5 节),或镜像名不符合约定(必须是 `pi-coop` / `pigo-coop` / `claude-coop`) | +| run_coop 容器秒退 | 检查 `MODEL` / `BASE_URL` / `API_KEY` 是否有效;`docker logs <容器名>` 看具体报错 | +| 长命令超时 | `AGENT_TOOL_TIMEOUT_SECONDS` / `AGENT_REQUEST_TIMEOUT_SECONDS` 调大(默认 600s) | +| Web 界面 401 | 设置了 `AGENT_AUTH_CODE`,先在弹窗输入授权码 | +| 历史压缩后丢上下文 | 调大 `AGENT_KEEP_RECENT_MESSAGES`(默认 50)与 `AGENT_COMPACTION_TOKENS`(默认 16384) | + +### 7.3 安全检查清单 + +- [ ] `.env` 与 API Key 未提交到 git(`.gitignore` 已包含 `.env`) +- [ ] 云端公开部署已设置 `AGENT_AUTH_CODE` +- [ ] API Key 泄露后立即在供应商侧吊销并轮换 +- [ ] Docker Socket 未暴露给不受信任的网络 +- [ ] `AGENT_WORKSPACE` 限定在专用目录,避免指向系统根目录 + +--- + +## 8. 端口与产物速查 + +| 场景 | 端口 | 产物 | +| --- | --- | --- | +| 本地运行 | 8080(`AGENT_PORT`) | — | +| pi/pigo/claude 协作镜像 | — | `pi-coop` / `pigo-coop` / `claude-coop` 镜像 | +| Tsecbench 托管 | 8080(容器内) | `tencent/agent.tar.gz` | diff --git a/README.md b/README.md new file mode 100644 index 0000000..685dc4b --- /dev/null +++ b/README.md @@ -0,0 +1,40 @@ +# Local Agent + +基于 Go(Gin)的多轮对话 AI Agent,ChatGPT 风格 Web 界面,可通过真实工具(Shell、文件、Python、Docker、协作容器)自主完成复杂任务。 + +## 核心特性 + +- **多轮工具调用**:`run_bash` / `read_file` / `write_file` / `run_python` / `list_directory` +- **Docker 操作**:`docker_ps` / `docker_exec` / `docker_start` 等,支持远程 Docker Socket +- **协作任务**:`run_coop` 调起独立协作容器(pi / pigo / claude 三种引擎),异步完成大型任务 +- **流式输出**:SSE 逐 token 输出,超长上下文自动压缩 +- **多协议 LLM**:OpenAI / Anthropic 兼容接口,Web 界面可切换模型与密钥 + +## 快速开始 + +```powershell +go run . +``` + +打开 http://localhost:8080 ,在 Web 设置中填入 API Key 即可使用。 + +## 部署 + +支持三种方式:本地运行、Docker 协作引擎、Tsecbench 托管镜像。 + +**详细步骤见 [DEPLOY.md](./DEPLOY.md)**(含全部环境变量说明、镜像构建、托管模式与排障)。 + +## 目录结构 + +| 路径 | 说明 | +| --- | --- | +| `main.go` / `internal/` | Go 控制层(Agent 核心 + HTTP/WS 服务) | +| `web/static/` | 前端静态资源 | +| `pi-coop/` | pi 协作引擎运行时(默认引擎) | +| `pigo/` | pigo 协作引擎(Go 实现) | +| `claude code/` | Claude Code 协作引擎 | +| `tencent/` | Tsecbench 托管模式镜像构建 | + +## 安全提醒 + +本程序具备命令执行与 Docker 操作能力,云端部署务必设置 `AGENT_AUTH_CODE` 授权码;API Key 通过环境变量注入,切勿提交到仓库。 diff --git a/claude code/AGENTS.md b/claude code/AGENTS.md new file mode 100644 index 0000000..6c50937 --- /dev/null +++ b/claude code/AGENTS.md @@ -0,0 +1,37 @@ +# 单 agent 任务协议 + +你是独立的任务执行 agent,完成 /blackboard/task.md 中的任务。任务信息通过工作区与黑板目录交换。 + +## 目录结构 + +| 路径 | 用途 | +| --- | --- | +| /blackboard/task.md | 任务描述,每轮都要重读,勿修改 | +| /blackboard/workspace/ | 你的工作区(你的 cwd),所有产物写在这里 | +| /blackboard/CLAUDE.md | 本协议(副本在你的工作区,勿修改) | +| /blackboard/DONE | 完成标记,存在即表示任务已交付完成,只用 blackboard 工具 done 创建 | +| /blackboard/logs/ | 每轮运行日志(supervisor 维护) | +| /blackboard/sessions/ | 会话 ID(supervisor 维护,供跨轮恢复) | +| /blackboard/result.json | 任务结果(supervisor 结束时生成,勿手动修改) | + +## 工具说明 + +你有三个 blackboard MCP 工具(命名空间 mcp__blackboard__): + +- **read**(无参数或 path):读取黑板。无 path 返回全局快照(task.md / messages / workspace / DONE 状态);有 path 返回指定文件内容。 +- **post**(file, content):原子追加一条进度消息到 messages/。file 必须是裸 .md 名(如 round-1-a.md)。 +- **done**(summary):原子创建 DONE 标记,summary 写最终交付总结。只能调用一次。 + +## 每轮流程 + +1. **读取**:用 blackboard read 读取 /blackboard/task.md(或直接 read 无参看全局快照),明确任务与提交规则。 +2. **评估**:检查工作区已有产物,判断进度与缺口,不重复已完成工作。 +3. **执行**:推进任务——执行命令、读写文件、编写产物到 /blackboard/workspace/。 +4. **提交**:拿到 flag 的任务,立即按 task.md 中约定的规则提交,并把提交响应与得分写入产物。 +5. **判定**:全部工作完成、交付物完整后,用 blackboard done 创建 DONE(summary 为最终交付总结,含关键结果/flag/提交响应/产物清单)。DONE 只能创建一次;宁可多轮打磨,不要过早宣布完成;确认无法解出时同样用 done 标记并写明「未解出」与已尝试内容。 + +## 红线 + +- 不直接写 /blackboard/DONE——必须走 blackboard done 工具,保证原子创建。 +- 不伪造实验、命令或提交结果;如实记录。 +- 不修改 /blackboard/task.md、/blackboard/CLAUDE.md 与 /blackboard/result.json。 diff --git a/claude code/Dockerfile b/claude code/Dockerfile new file mode 100644 index 0000000..bb22a1f --- /dev/null +++ b/claude code/Dockerfile @@ -0,0 +1,75 @@ +# Claude Code 协作镜像(单 agent 版,与 blackboard 协议适配) +# +# 构建(在仓库根目录执行,构建上下文为 claude code/ 目录): +# cd /mnt/e/Code/Go/awesomeProject/agent +# DOCKER_BUILDKIT=0 docker build -f "claude code/Dockerfile" -t claude-coop "claude code/" +# +# 基础镜像用 node:22-alpine(自带 node 22 + npm,满足 Claude Code 的 Node 要求)。 +# 在其上装系统工具(curl/wget/git/jq/openssl/ripgrep)和 python3 + 常用库, +# 再全局安装 @anthropic-ai/claude-code 和 MCP SDK,最后放入协作组件 +#(blackboard MCP server / supervisor / prompts)。 +# +# 与 pi-coop 的差异: +# - 用 claude-code 替代 pi-coding-agent +# - blackboard 工具从 pi extension 改为独立 MCP server(stdio 协议) +# - 模型配置走 ANTHROPIC_* 环境变量(不是 pi provider 体系) + +FROM node:22-alpine + +# ---- 系统工具 + Python ---- +# 先把 apk 源换成清华镜像并用 http(node:22-alpine 的 Alpine 3.24 用官方 CDN +# 拉 APKINDEX 会报 TLS: unspecified error,改 http 绕过;装上 ca-certificates 后 +# TLS 即恢复正常)。之后 apk add 装系统工具 + python3。 +RUN sed -i 's|https://dl-cdn.alpinelinux.org|http://mirrors.tuna.tsinghua.edu.cn|g' /etc/apk/repositories \ + && apk add --no-cache \ + bash ca-certificates git ripgrep \ + curl wget jq openssl file \ + python3 py3-pip + +# ---- Python 常用库 ---- +# 用清华 PyPI 镜像加速;--no-cache-dir 避免膨胀;--break-system-packages 放行 Alpine pip 全局安装。 +RUN pip3 install --no-cache-dir --break-system-packages \ + -i https://pypi.tuna.tsinghua.edu.cn/simple \ + requests urllib3 certifi idna charset-normalizer \ + cryptography pyOpenSSL paramiko pyjwt \ + PyYAML beautifulsoup4 lxml \ + python-dotenv click tqdm \ + numpy pandas \ + httpx aiohttp websockets dnspython \ + psutil pillow + +# ---- 全局安装 Claude Code ---- +# 使用 npmmirror 镜像加速国内下载。 +# @anthropic-ai/claude-code:提供 claude CLI(headless 模式入口) +RUN npm install -g \ + --registry=https://registry.npmmirror.com \ + @anthropic-ai/claude-code@latest + +# ---- 协作组件 ---- +# blackboard MCP server 放 /mcp,supervisor 与 prompts 放默认路径。 +RUN mkdir -p /mcp /prompts /blackboard +COPY blackboard-server.mjs /mcp/blackboard-server.mjs +COPY mcp-servers.json /mcp-servers.json +COPY supervisor.sh /usr/local/bin/supervisor.sh +COPY AGENTS.md /prompts/AGENTS.md +COPY prompts /prompts +RUN chmod +x /usr/local/bin/supervisor.sh + +# ---- MCP SDK 本地安装 ---- +# ESM 模块解析要求依赖在脚本所在目录的 node_modules 内(全局安装的包不会被 +# ESM import 解析到,NODE_PATH 也只对 CommonJS require 生效)。因此在 /mcp +# 目录本地安装 @modelcontextprotocol/sdk,让 blackboard-server.mjs 能 import 到。 +RUN cd /mcp && npm init -y >/dev/null 2>&1 \ + && npm install \ + --registry=https://registry.npmmirror.com \ + @modelcontextprotocol/sdk@latest + +# ---- 非 root 用户 ---- +# node:22-alpine 自带 node 用户(uid 1000),直接复用,与 pi-coop 的 uid 1000 对齐, +# 便于 blackboard 挂载卷权限一致。 +# 注意:--dangerously-skip-permissions 不能以 root 运行,必须用非 root 用户。 +RUN chown -R node:node /mcp /prompts /blackboard /mcp-servers.json +USER node +WORKDIR /blackboard + +ENTRYPOINT ["/usr/local/bin/supervisor.sh"] diff --git a/claude code/RUN.md b/claude code/RUN.md new file mode 100644 index 0000000..809d9d7 --- /dev/null +++ b/claude code/RUN.md @@ -0,0 +1,70 @@ +# claude-coop 镜像构建与运行说明 + +## 构建镜像 + +在仓库根目录(`agent/`)执行: + +```bash +DOCKER_BUILDKIT=0 docker build -f "claude code/Dockerfile" -t claude-coop "claude code/" +``` + +> 构建上下文是 `claude code/` 目录(注意路径含空格,需引号包裹)。 + +## 镜像内容 + +| 组件 | 说明 | +| --- | --- | +| 基础镜像 | `node:22-alpine`(node 22 + npm) | +| 编码 agent | `@anthropic-ai/claude-code`(claude CLI,headless 模式入口) | +| MCP runtime | `@modelcontextprotocol/sdk`(blackboard MCP server 依赖) | +| 系统工具 | curl wget git jq openssl ripgrep file | +| Python | python3 + 常用库(requests/cryptography/paramiko/PyYAML/numpy 等) | +| blackboard 工具 | MCP server(`/mcp/blackboard-server.mjs`),提供 read/post/done 三个工具 | + +## 与 pi-coop 的差异 + +| 维度 | pi-coop | claude-coop | +| --- | --- | --- | +| 编码 agent | pi(@earendil-works/pi-coding-agent) | Claude Code(@anthropic-ai/claude-code) | +| 工具注册 | pi extension(-e flag) | MCP server(--mcp-config) | +| 模型配置 | provider 体系(--provider coop --model X) | ANTHROPIC_* 环境变量 | +| 协议支持 | openai + anthropic | 仅 anthropic(Claude Code 限制) | +| session 提取 | NDJSON 首行 session header | JSON result 的 session_id 字段 | +| session 恢复 | --session-id | --resume | +| 权限 | 无(pi 无权限弹窗) | --dangerously-skip-permissions(headless 必需,非 root) | + +## 环境变量契约(与 pi-coop 一致) + +| 变量 | 说明 | 必填 | +| --- | --- | --- | +| MODEL | 模型名(如 deepseek-v4-flash) | 是 | +| BASE_URL | Anthropic 兼容 API base-url(不带 /v1) | 是 | +| API_KEY | API key | 是 | +| PROTOCOL | 协议(claude-coop 仅支持 anthropic) | 否 | +| TASK | 任务描述(或挂载 $BB/task.md) | 是 | +| BLACKBOARD | 黑板目录(默认 /blackboard) | 否 | +| ROUND_MAX | 最大轮次(默认 10) | 否 | +| TIMEOUT | 单轮超时秒数(默认 600) | 否 | +| FAIL_MODE | agent 失败时:stop(默认)/ continue | 否 | + +## 手动测试 + +```bash +# 测试镜像工具链 +docker run --rm --entrypoint sh claude-coop -c 'which claude node python3 curl wget jq; claude --version' + +# 测试 blackboard MCP server +docker run --rm --entrypoint sh claude-coop -c 'BB=/tmp/bb mkdir -p /tmp/bb && BB=/tmp/bb node /mcp/blackboard-server.mjs & sleep 1 && kill %1' +``` + +## 控制层适配 + +控制层 `tools_coop.go` 中 `coopImage` 常量当前为 `"pi-coop"`。如需切换到 claude-coop: + +```go +const coopImage = "claude-coop" +``` + +环境变量注入逻辑无需改动(MODEL/BASE_URL/API_KEY/PROTOCOL 契约一致),但需注意: +- Claude Code 仅支持 Anthropic 协议端点,`coopBaseURL` 对 anthropic 不再补 `/v1`(已修复) +- `PROTOCOL=anthropic` 时,`BASE_URL` 应为 `https://api.deepseek.com/anthropic` diff --git a/claude code/blackboard-server.mjs b/claude code/blackboard-server.mjs new file mode 100644 index 0000000..56ede41 --- /dev/null +++ b/claude code/blackboard-server.mjs @@ -0,0 +1,297 @@ +#!/usr/bin/env node +/** + * blackboard MCP server:让 Claude Code 适配 blackboard 协作协议。 + * + * 与 pi-coop/coop.ts 的 blackboard 工具行为对齐: + * - read 无 path → 全局快照(task.md / messages / workspace / DONE 状态) + * 有 path → 单文件内容 + * - post 原子追加一条消息到 messages/(O_APPEND,POSIX 下并发不交错) + * - done 原子创建 DONE 标记(O_CREAT|O_EXCL,保证唯一创建) + * + * 路径安全:所有传入路径都做 traversal 校验,拒绝绝对路径与 .. 越界。 + * + * 通信协议:MCP stdio(JSON-RPC 2.0)。Claude Code 通过 --mcp-config 加载本 server, + * 工具命名空间为 mcp__blackboard__read / mcp__blackboard__post / mcp__blackboard__done。 + * + * 环境变量: + * BB 黑板根目录(默认 /blackboard,由 supervisor 注入) + */ +import { Server } from "@modelcontextprotocol/sdk/server/index.js"; +import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; +import { + CallToolRequestSchema, + ListToolsRequestSchema, +} from "@modelcontextprotocol/sdk/types.js"; +import { + openSync, + closeSync, + writeFileSync, + appendFileSync, + readFileSync, + readdirSync, + statSync, + mkdirSync, + existsSync, +} from "node:fs"; +import { join, resolve, sep, isAbsolute, normalize, basename } from "node:path"; + +// clean 等价于 normalize:规范化路径(消除 . 和 ..,合并分隔符)。 +// pi 的 coop.ts 用了内部 clean,Node 标准 path 模块对应方法是 normalize。 +const clean = normalize; + +// 单条 post 上限:保证单次原子写,且不撑爆上下文(与 pigo / pi-coop 一致)。 +const MAX_POST_BYTES = 32 * 1024; +// 单文件 read 上限。 +const MAX_READ_BYTES = 32 * 1024; + +const bbRoot = process.env.BB ?? "/blackboard"; + +// ---- 路径安全:safePath 把相对路径解析到 blackboard 根内,拒绝绝对路径与 .. 越界 ---- +function safePath(p) { + const trimmed = (p ?? "").trim(); + if (trimmed === "") throw new Error("empty path"); + const c = clean(trimmed); + if (isAbsolute(c)) throw new Error(`path "${p}" must be relative to the blackboard root`); + const rootClean = clean(bbRoot); + const full = join(rootClean, c); + if (full !== rootClean && !full.startsWith(rootClean + sep)) { + throw new Error(`path "${p}" escapes the blackboard root`); + } + return full; +} + +function truncate(s, max) { + return s.length > max ? s.slice(0, max) + "\n..." : s; +} + +// ---- blackboard 工具实现 ---- + +// read:无 path → 全局快照;有 path → 单文件内容。 +function readAction(pathParam) { + const p = pathParam != null ? String(pathParam) : ""; + if (p.trim() !== "") return readFile(p); + + const b = []; + // task.md + try { + const data = readFileSync(join(bbRoot, "task.md"), "utf8"); + b.push("# task.md\n" + truncate(data, MAX_READ_BYTES) + "\n"); + } catch { + b.push("# task.md\n\n"); + } + + // messages/ + b.push("\n# messages/ (" + listNames("messages").join(", ") + ")\n"); + for (const name of listListing("messages")) b.push(" - " + name + "\n"); + + // workspace/ + b.push("\n# workspace/ (your workspace)\n"); + for (const name of listListing("workspace")) b.push(" - " + name + "\n"); + + // DONE + let doneText = ""; + try { + doneText = truncate(readFileSync(join(bbRoot, "DONE"), "utf8"), 4096); + } catch {} + b.push("\n# DONE\n"); + b.push(doneText ? "EXISTS:\n" + doneText + "\n" : "\n"); + + // environment + const env = []; + for (const k of ["BB", "ROUND", "NAME"]) { + const v = (process.env[k] ?? "").trim(); + if (v) env.push(` ${k}=${v}`); + } + if (env.length) b.push("\n# environment\n" + env.join("\n") + "\n"); + + return b.join(""); +} + +function readFile(p) { + const full = safePath(p); + let info; + try { + info = statSync(full); + } catch (e) { + throw new Error(`blackboard read: ${p}: ${e instanceof Error ? e.message : String(e)}`); + } + if (info.isDirectory()) throw new Error(`blackboard read: ${p} is a directory; only files can be read`); + const data = readFileSync(full, "utf8"); + return "# " + p + "\n" + truncate(data, MAX_READ_BYTES); +} + +// post:原子追加一条消息到 messages/。file 必须是裸 *.md 名。 +function postAction(fileParam, contentParam) { + const name = String(fileParam ?? "").trim(); + if (!validMessageName(name)) { + throw new Error('blackboard post: file must be a bare name ending in .md (e.g. "round-1-a.md"), no path separators, no ".."'); + } + const content = String(contentParam ?? "").trim(); + if (content === "") throw new Error("blackboard post: content must not be empty"); + if (Buffer.byteLength(content) > MAX_POST_BYTES) { + throw new Error(`blackboard post: content too large (${Buffer.byteLength(content)} bytes, max ${MAX_POST_BYTES})`); + } + const dir = join(bbRoot, "messages"); + mkdirSync(dir, { recursive: true }); + // O_APPEND 单次写:POSIX 下并发 post 不会交错字节。 + appendFileSync(join(dir, name), content + "\n"); + return `Message appended to messages/${name}`; +} + +// done:原子创建 DONE 标记。'wx' = O_CREAT|O_EXCL|O_WRONLY,保证唯一创建。 +function doneAction(summaryParam) { + const summary = String(summaryParam ?? "").trim(); + if (summary === "") throw new Error("blackboard done: summary must not be empty (include the final delivery summary)"); + const header = "Blackboard cooperation DONE\ncreated: " + new Date().toISOString() + "\n\n"; + const target = join(bbRoot, "DONE"); + let fd; + try { + fd = openSync(target, "wx"); + } catch (e) { + if (existsSync(target)) { + const existing = readFileSync(target, "utf8"); + throw new Error("blackboard done: DONE already exists — cooperation already finished:\n" + truncate(existing, 4096)); + } + throw e; + } + try { + writeFileSync(fd, header + summary); + } finally { + closeSync(fd); + } + return "DONE marker created. Cooperation finished."; +} + +// 校验 post 文件名:裸 *.md,无分隔符,无 . / .. +function validMessageName(name) { + if (name === "" || !name.endsWith(".md")) return false; + if (/[\\/]/.test(name) || name === "." || name === "..") return false; + const base = name.slice(0, -3); + if (base === "" || base.startsWith(".") || base.includes("..")) return false; + return true; +} + +function listNames(sub) { + try { + return readdirSync(join(bbRoot, sub)).sort(); + } catch { + return []; + } +} + +function listListing(sub) { + try { + return readdirSync(join(bbRoot, sub)) + .sort() + .map((name) => { + try { + const info = statSync(join(bbRoot, sub, name)); + if (info.isDirectory()) return `${name} (dir)`; + return `${name} (${info.size} bytes, ${info.mtime.toISOString().slice(11, 19)})`; + } catch { + return name; + } + }); + } catch { + return []; + } +} + +// ---- MCP server ---- +const server = new Server( + { name: "blackboard", version: "1.0.0" }, + { capabilities: { tools: {} } }, +); + +// 列出工具:三个 action 拆成独立工具,便于 Claude Code 自动批准白名单匹配。 +server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: [ + { + name: "read", + description: + "Read the task blackboard used by the coop runner. " + + "Without a path: returns a global snapshot (task.md, messages listing, workspace listing, DONE state). " + + "With a path like \"workspace/exploit.py\": returns the contents of that one file. " + + "The blackboard root, current round and your name are in the environment as BB, ROUND, NAME.", + inputSchema: { + type: "object", + properties: { + path: { + type: "string", + description: 'A path under the blackboard root to fetch (e.g. messages/round-1-a.md). Omit for the global snapshot.', + }, + }, + }, + }, + { + name: "post", + description: + "Atomically append a progress note to the blackboard messages/ directory. " + + "The file must be a bare .md name under messages/ (e.g. round-1-a.md). Max 32 KiB per post.", + inputSchema: { + type: "object", + properties: { + file: { + type: "string", + description: 'Bare file name under messages/, must end in .md (e.g. round-1-a.md).', + }, + content: { + type: "string", + description: "The message body (max 32 KiB).", + }, + }, + required: ["file", "content"], + }, + }, + { + name: "done", + description: + "Atomically create the DONE marker with a final delivery summary. " + + "Only call this when the deliverable is truly complete (or confirmed unsolvable). Fails if DONE already exists.", + inputSchema: { + type: "object", + properties: { + summary: { + type: "string", + description: "Final delivery summary written into DONE (include key results / flag / submit response / artifact list).", + }, + }, + required: ["summary"], + }, + }, + ], +})); + +// 调用工具:根据 name 分发到对应实现,错误转为文本返回(不抛异常,保持 MCP 协议干净)。 +server.setRequestHandler(CallToolRequestSchema, async (request) => { + const { name, arguments: args } = request.params; + try { + let text; + switch (name) { + case "read": + text = readAction(args?.path); + break; + case "post": + text = postAction(args?.file, args?.content); + break; + case "done": + text = doneAction(args?.summary); + break; + default: + return { + content: [{ type: "text", text: `blackboard: unknown tool "${name}" (want read|post|done)` }], + isError: true, + }; + } + return { content: [{ type: "text", text }] }; + } catch (e) { + return { + content: [{ type: "text", text: `blackboard ${name}: ${e instanceof Error ? e.message : String(e)}` }], + isError: true, + }; + } +}); + +// 启动 stdio transport,由 Claude Code 通过子进程方式拉起。 +const transport = new StdioServerTransport(); +await server.connect(transport); diff --git a/claude code/mcp-servers.json b/claude code/mcp-servers.json new file mode 100644 index 0000000..2fc715c --- /dev/null +++ b/claude code/mcp-servers.json @@ -0,0 +1,9 @@ +{ + "mcpServers": { + "blackboard": { + "command": "node", + "args": ["/mcp/blackboard-server.mjs"], + "env": {} + } + } +} diff --git a/claude code/prompts/agent.md b/claude code/prompts/agent.md new file mode 100644 index 0000000..75f44ff --- /dev/null +++ b/claude code/prompts/agent.md @@ -0,0 +1,18 @@ +# 角色:agent(主执行者) + +你是 Claude Code 单 agent 任务执行者,独立完成 /blackboard/task.md 中的任务。 + +## 你的职责 +- 理解任务、制定方案、完成核心产出。 +- 所有产物写入你的工作区(你的 cwd,即 /blackboard/workspace/),按需命名。 +- 需要提交 flag 的任务:拿到 flag 后立即按任务中的提交规则提交,并把提交响应与得分记录在产物中。 + +## 工作方式 +- 每轮:用 blackboard read 读取 /blackboard/task.md 与工作区已有产物 → 推进任务 → 将进展与产物写入工作区。 +- 环境变量 ROUND、NAME、BB 由 supervisor 提供;需要记录进展时用 blackboard post,创建完成标记时用 blackboard done(BB 指向黑板根目录)。 +- 不要伪造命令结果、文件内容或提交响应;如实记录。 +- 全部工作完成、交付物完整时,用 blackboard done 创建完成标记(summary 写最终交付总结,含关键结果、flag、提交响应、产物清单)。 + +## 注意 +- 若本轮未完成,下一轮会用 --resume 恢复你的会话继续推进,跨轮保持上下文。 +- 宁可多轮打磨,不要过早宣布完成;确认无法解出时同样用 blackboard done 标记,summary 写明「未解出」与已尝试内容,让调度方及时切换下一题。 diff --git a/claude code/supervisor.sh b/claude code/supervisor.sh new file mode 100644 index 0000000..8d5ea98 --- /dev/null +++ b/claude code/supervisor.sh @@ -0,0 +1,215 @@ +#!/usr/bin/env bash +# Claude Code 协作 supervisor(单 agent 版) +# +# 在容器内运行一个 Claude Code headless 进程(claude -p)完成 /blackboard/task.md 中的任务。 +# agent 的工作区为 /blackboard/workspace,产物与轮次日志保留在黑板上; +# 任务完成(agent 用 blackboard MCP 工具创建 DONE)后,supervisor 把结果 +# 以结构化 JSON(result.json)落盘并打印一行摘要,供外部调度方直接解析。 +# +# 与 pi-coop supervisor 的关键差异(Claude Code CLI 与 pi 不同): +# 1. 模型接入:Claude Code 用 ANTHROPIC_BASE_URL / ANTHROPIC_AUTH_TOKEN / +# ANTHROPIC_MODEL 环境变量(不是 --provider/--model flag)。 +# 第三方 Anthropic 兼容端点(如 DeepSeek)需 ANTHROPIC_API_KEY="" 触发 +# 回退到 ANTHROPIC_AUTH_TOKEN,且 ANTHROPIC_BASE_URL 不带 /v1(SDK 自己追加)。 +# 2. 工具注册:Claude Code 用 MCP server(--mcp-config),blackboard 工具 +# 通过 blackboard-server.mjs 提供,工具命名空间为 mcp__blackboard__*。 +# 3. 输出格式:claude -p --output-format json 输出单个 JSON,含 session_id +# 字段,直接用 jq 提取即可(无需像 pi 那样解析 NDJSON 首行)。 +# 4. session 恢复:claude -p --resume ,跨轮保持上下文。 +# 5. 权限:--dangerously-skip-permissions 跳过所有权限弹窗(headless 必需), +# 且必须以非 root 用户运行(node:22-alpine 自带 node 用户,uid 1000)。 +# +# 环境变量(均由外部调用方传入,与 pi-coop 完全一致): +# MODEL 模型名,如 deepseek-v4-flash (必填) +# BASE_URL Anthropic 兼容 API base-url (必填) +# API_KEY API key(注入为 ANTHROPIC_AUTH_TOKEN) (必填) +# PROTOCOL 协议(仅 anthropic 受 Claude Code 支持) (可选) +# TASK 任务描述 (必填;或挂载 $BB/task.md) +# BLACKBOARD 黑板目录 (默认 /blackboard) +# ROUND_MAX 最大轮次 (默认 10) +# TIMEOUT 单轮超时秒数,0=不超时 (默认 600) +# FAIL_MODE agent 失败时:stop=立即退出(默认)| continue=继续下一轮 +set -u + +BB="${BLACKBOARD:-/blackboard}" +PROMPTS="${PROMPTS:-/prompts}" +MCP_CONFIG="${MCP_CONFIG:-/mcp-servers.json}" +ROUND_MAX="${ROUND_MAX:-10}" +TIMEOUT="${TIMEOUT:-600}" +FAIL_MODE="${FAIL_MODE:-stop}" + +log() { echo "[supervisor] $*"; } + +# emit_result 把任务结果以结构化 JSON 写入 $BB/result.json 并打印一行摘要。 +# status: solved | unsolved | error | timeout +emit_result() { + local status="$1" summary="$2" code="$3" + local flag="" artifacts_json="[]" summary_json='""' flag_json='""' + + # flag 优先从 summary 提取,其次在 workspace 产物中全量检索 + if [ -n "$summary" ]; then + flag=$(printf '%s' "$summary" | grep -oE 'flag\{[^}]+\}' | head -1) + fi + if [ -z "$flag" ] && [ -d "$BB/workspace" ]; then + flag=$(grep -rhoE 'flag\{[^}]+\}' "$BB/workspace" 2>/dev/null | head -1) + fi + + # 产物清单:workspace 下的全部文件(最多 50 个),经 python3 转义为 JSON 数组 + if [ -d "$BB/workspace" ]; then + artifacts_json=$( + cd "$BB/workspace" && find . -type f 2>/dev/null | sed 's|^\./||' | head -50 | + python3 -c 'import json,sys; print(json.dumps([l.rstrip("\n") for l in sys.stdin]))' 2>/dev/null + ) + [ -z "$artifacts_json" ] && artifacts_json="[]" + fi + + # summary 截断到 4000 字符并经 python3 转义,防止引号/换行破坏 JSON + summary_json=$(printf '%s' "$summary" | head -c 4000 | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))' 2>/dev/null) + [ -z "$summary_json" ] && summary_json='""' + flag_json=$(printf '%s' "$flag" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))' 2>/dev/null) + [ -z "$flag_json" ] && flag_json='""' + + cat > "$BB/result.json" <&2 + exit 2 +fi +if [ -z "${TASK:-}" ] && [ ! -f "$BB/task.md" ]; then + echo "错误:请通过 TASK 环境变量或挂载 $BB/task.md 提供任务" >&2 + exit 2 +fi +if [ -f "$BB/DONE" ]; then + log "黑板已有完成标记,如需重新开始请删除 $BB/DONE" + exit 0 +fi + +# ---- 初始化黑板 ---- +mkdir -p "$BB/workspace" "$BB/logs" "$BB/sessions" +if [ -n "${TASK:-}" ]; then + printf '%s\n' "$TASK" > "$BB/task.md" +fi +# 协议注入:agent 的 cwd 是 $BB/workspace,AGENTS.md 注入链从 cwd 起, +# Claude Code 会自动加载 cwd 下的 CLAUDE.md(AGENTS.md 兼容),因此把协议 +# 副本放进工作区(勿修改,它是每轮系统提示的一部分) +cp "$PROMPTS/AGENTS.md" "$BB/workspace/CLAUDE.md" +cp "$PROMPTS/AGENTS.md" "$BB/workspace/AGENTS.md" +log "黑板初始化完成:$BB" +log "任务:$(head -c 200 "$BB/task.md")" + +# ---- 导出 Claude Code 模型配置环境变量 ---- +# 关键:ANTHROPIC_API_KEY 必须设为空字符串(而非 unset),Claude Code 才会 +# 回退到 ANTHROPIC_AUTH_TOKEN。BASE_URL 不带 /v1(SDK 自动追加 /v1/messages)。 +export ANTHROPIC_BASE_URL="$BASE_URL" +export ANTHROPIC_AUTH_TOKEN="$API_KEY" +export ANTHROPIC_API_KEY="" +export ANTHROPIC_MODEL="$MODEL" +# 把 opus/sonnet/haiku 三个别名都映射到目标模型,避免别名解析失败。 +export ANTHROPIC_DEFAULT_OPUS_MODEL="$MODEL" +export ANTHROPIC_DEFAULT_SONNET_MODEL="$MODEL" +export ANTHROPIC_DEFAULT_HAIKU_MODEL="$MODEL" +# 禁用非必要流量(遥测/登录检测),第三方端点必需。 +export CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 + +# ---- 单轮运行:claude -p --output-format json,从 JSON 提取 session_id ---- +run_agent() { + local round="$1" + local session_file="$BB/sessions/agent.session" + local task_prompt="轮次 $round 开始。先读取 /blackboard/task.md 中的任务,检查工作区已有产物与 DONE 状态,然后继续推进任务。全部工作完成、交付物完整时,用 blackboard 工具 done 创建完成标记(summary 写最终交付总结,含关键结果/flag/提交响应/产物清单)。" + local args=(-p "$task_prompt" + --output-format json + --mcp-config "$MCP_CONFIG" + --allowedTools "mcp__blackboard__*,Read,Write,Edit,Bash" + --append-system-prompt-file "$PROMPTS/agent.md" + --dangerously-skip-permissions) + if [ -s "$session_file" ]; then + args+=(--resume "$(cat "$session_file")") + fi + + export ROUND="$round" NAME="agent" BB="$BB" + local logfile="$BB/logs/round-$round.log" + log "第 $round 轮开始运行 agent(日志:$logfile)" + # Claude Code 用 --cwd / 进程 cwd 作工作区,先进入 $BB/workspace + if [ "$TIMEOUT" -gt 0 ] 2>/dev/null; then + (cd "$BB/workspace" && timeout "$TIMEOUT" claude "${args[@]}") > "$logfile" 2>&1 + else + (cd "$BB/workspace" && claude "${args[@]}") > "$logfile" 2>&1 + fi + local rc=$? + + # 从 --output-format json 的结果提取 session_id 供下一轮 --resume + # JSON 形如 {"type":"result","session_id":"",...} + local sid + sid=$(python3 -c ' +import json, sys +try: + d = json.loads(sys.stdin.read()) + if d.get("type") == "result": + print(d.get("session_id", "")) +except Exception: + pass +' < "$logfile" 2>/dev/null) + if [ -n "$sid" ]; then + printf '%s' "$sid" > "$session_file" + fi + return $rc +} + +# ---- 主循环:逐轮运行 agent,直到 DONE 或达到轮次上限 ---- +for round in $(seq 1 "$ROUND_MAX"); do + [ -f "$BB/DONE" ] && break + log "===== 第 $round 轮开始 =====" + + run_agent "$round" + rc=$? + + if [ "$rc" -eq 0 ]; then + log "第 $round 轮完成" + else + log "第 $round 轮失败(exit=$rc),日志见 $BB/logs/round-$round.log" + if [ "$FAIL_MODE" = "stop" ]; then + # 兜底:agent 虽超时/失败,但黑板已有提交成功证据(workspace 内 *.md 含 correct:true) + # → 视为完成正常退出,避免"flag 已提交却因未写 DONE 被强杀(exit 143)"。 + _done=0 + for _f in "$BB"/workspace/*.md "$BB"/messages/*.md; do + [ -f "$_f" ] && grep -q '"correct":true' "$_f" && _done=1 && break + done + if [ "$_done" = "1" ]; then + log "检测到提交成功证据,自动标记完成" + { echo "Task finished (auto-detected submit success)"; } > "$BB/DONE" + emit_result solved "$(cat "$BB/DONE")" 0 + exit 0 + fi + if [ "$rc" -eq 124 ]; then + emit_result timeout "" "$rc" + else + emit_result error "agent 退出码 $rc,日志见 logs/round-$round.log" "$rc" + fi + exit "$rc" + fi + fi +done + +if [ -f "$BB/DONE" ]; then + log "检测到完成标记,任务结束" + emit_result solved "$(cat "$BB/DONE")" 0 + exit 0 +fi + +log "达到最大轮次 $ROUND_MAX 仍未完成,请检查黑板产物与日志" +emit_result unsolved "" 1 +exit 1 diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..dd47a54 --- /dev/null +++ b/go.mod @@ -0,0 +1,64 @@ +module agent + +go 1.25.0 + +require ( + github.com/docker/docker v28.5.2+incompatible + github.com/gin-gonic/gin v1.10.0 +) + +require ( + github.com/Microsoft/go-winio v0.6.2 // indirect + github.com/bytedance/sonic v1.11.6 // indirect + github.com/bytedance/sonic/loader v0.1.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/cloudwego/base64x v0.1.4 // indirect + github.com/cloudwego/iasm v0.2.0 // indirect + github.com/containerd/errdefs v1.0.0 // indirect + github.com/containerd/errdefs/pkg v0.3.0 // indirect + github.com/containerd/log v0.1.0 // indirect + github.com/distribution/reference v0.6.0 // indirect + github.com/docker/go-connections v0.8.1 // indirect + github.com/docker/go-units v0.5.0 // indirect + github.com/felixge/httpsnoop v1.1.0 // indirect + github.com/gabriel-vasile/mimetype v1.4.3 // indirect + github.com/gin-contrib/sse v0.1.0 // indirect + github.com/go-logr/logr v1.4.4 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/go-playground/validator/v10 v10.20.0 // indirect + github.com/goccy/go-json v0.10.2 // indirect + github.com/gorilla/websocket v1.5.3 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.2.7 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/moby/docker-image-spec v1.3.1 // indirect + github.com/moby/sys/atomicwriter v0.1.0 // indirect + github.com/moby/term v0.5.2 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/morikuni/aec v1.1.0 // indirect + github.com/opencontainers/go-digest v1.0.0 // indirect + github.com/opencontainers/image-spec v1.1.1 // indirect + github.com/pelletier/go-toml/v2 v2.2.2 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/ugorji/go/codec v1.2.12 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.70.0 // indirect + go.opentelemetry.io/otel v1.45.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 // indirect + go.opentelemetry.io/otel/metric v1.45.0 // indirect + go.opentelemetry.io/otel/trace v1.45.0 // indirect + golang.org/x/arch v0.8.0 // indirect + golang.org/x/crypto v0.54.0 // indirect + golang.org/x/net v0.57.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/text v0.40.0 // indirect + golang.org/x/time v0.15.0 // indirect + google.golang.org/protobuf v1.36.11 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect + gotest.tools/v3 v3.5.2 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..7afcd73 --- /dev/null +++ b/go.sum @@ -0,0 +1,178 @@ +github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c h1:udKWzYgxTojEKWjV8V+WSxDXJ4NFATAsZjh8iIbsQIg= +github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= +github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= +github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= +github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc0= +github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4= +github.com/bytedance/sonic/loader v0.1.1 h1:c+e5Pt1k/cy5wMveRDyk2X4B9hF4g7an8N3zCYjJFNM= +github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= +github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= +github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y= +github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= +github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg= +github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= +github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= +github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= +github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= +github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk= +github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I= +github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= +github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= +github.com/docker/docker v28.5.2+incompatible h1:DBX0Y0zAjZbSrm1uzOkdr1onVghKaftjlSWt4AFexzM= +github.com/docker/docker v28.5.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= +github.com/docker/go-connections v0.8.1 h1:JibmG5hULs5qXSr/cp/w3Pw5fZuStt4MOHMUExb29/M= +github.com/docker/go-connections v0.8.1/go.mod h1:no1qkHdjq7kLMGUXYAduOhYPSJxxvgWBh7ogVvptn3Q= +github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= +github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= +github.com/felixge/httpsnoop v1.1.0 h1:3YtUj32ZZkqZtt3sZZsClsymw/QDuVfpNhoA31zeORc= +github.com/felixge/httpsnoop v1.1.0/go.mod h1:Zqxgdd+1Rkcz8euOqdr7lqgCRJztwr5hp9vDSi5UZCE= +github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= +github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= +github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE= +github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= +github.com/gin-gonic/gin v1.10.0 h1:nTuyha1TYqgedzytsKYqna+DfLos46nTv2ygFy86HFU= +github.com/gin-gonic/gin v1.10.0/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= +github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.20.0 h1:K9ISHbSaI0lyB2eWMPJo+kOS/FBExVwjEviJTixqxL8= +github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM= +github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= +github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF28c5ZQfqCBQ5g2xfk= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= +github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM= +github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= +github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= +github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= +github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo= +github.com/moby/sys/atomicwriter v0.1.0 h1:kw5D/EqkBwsBFi0ss9v1VG3wIkVhzGvLklJ+w3A14Sw= +github.com/moby/sys/atomicwriter v0.1.0/go.mod h1:Ul8oqv2ZMNHOceF643P6FKPXeCmYtlQMvpizfsSoaWs= +github.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU= +github.com/moby/sys/sequential v0.6.0/go.mod h1:uyv8EUTrca5PnDsdMGXhZe6CCe8U/UiTWd+lL+7b/Ko= +github.com/moby/term v0.5.2 h1:6qk3FJAFDs6i/q3W/pQ97SX192qKfZgGjCQqfCJkgzQ= +github.com/moby/term v0.5.2/go.mod h1:d3djjFCrjnB+fl8NJux+EJzu0msscUP+f8it8hPkFLc= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/morikuni/aec v1.1.0 h1:vBBl0pUnvi/Je71dsRrhMBtreIqNMYErSAbEeb8jrXQ= +github.com/morikuni/aec v1.1.0/go.mod h1:xDRgiq/iw5l+zkao76YTKzKttOp2cwPEne25HDkJnBw= +github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= +github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= +github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= +github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= +github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= +github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= +github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.70.0 h1:LMuyCAyfalSjDyjdC65nK6N0zoTT63+E/u95X0JovZI= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.70.0/go.mod h1:085m8qbm4hgc8rZWGDEa4vmyyo2c3nPxUslYUKUIU04= +go.opentelemetry.io/otel v1.45.0 h1:pdrWmLHofpubmArBv1LgFSv1Z0Ie/ppdZzu+kUN5EeU= +go.opentelemetry.io/otel v1.45.0/go.mod h1:XZxIqPapzEYnhNSScF5DIqXhm/rYi0FzCe2XddAwZfQ= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 h1:QRefszxJmfPdjXUUm3j6iDzY03mTPXMjqErFqQ67vUg= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0/go.mod h1:Tiz03lTBVBrm7eWZBOidzEaYaJa8tjwGUGv6d8mlTyk= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 h1:QBajQ2SrwQijzHyZbQlPsuIzpl/ll8DY6wPWsajeGcI= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0/go.mod h1:08ZQLjrPLQ6R4kAXvuOvODEer5Yh4CoFvll5qB2BCI8= +go.opentelemetry.io/otel/metric v1.45.0 h1:7Eg1uH7CJ5cXv9is6tnBe1FI6rj1nwUdbFypRm3br/M= +go.opentelemetry.io/otel/metric v1.45.0/go.mod h1:HAPbm1nd3p1PmFH7v2dR+6BjXxw+Lq4a2+pndMAm08s= +go.opentelemetry.io/otel/sdk v1.45.0 h1:4VVSMgQ83dUgW2aoX5f6JgLvHwIvzcuLnF9lUdCSpCw= +go.opentelemetry.io/otel/sdk v1.45.0/go.mod h1:Sr40LgXV7DsKMMJMKOhUWOgMWTfAaqvm2kF0g7ilwuA= +go.opentelemetry.io/otel/sdk/metric v1.45.0 h1:oVFszMfyj1Am6s24Vtc7wBb8BKLcwepJjNEYILuiE3o= +go.opentelemetry.io/otel/sdk/metric v1.45.0/go.mod h1:vUWUxDZvu1WVRj8JA8S0AdhsPrZoDpA2DdZauIh4mDA= +go.opentelemetry.io/otel/trace v1.45.0 h1:l/mP6Uv7oNO7/TblbhpbgMidxhq1uO/rPsikOyVhxag= +go.opentelemetry.io/otel/trace v1.45.0/go.mod h1:qoJJA2xNMnxRrdISU/kLtfUH2wNeQbiv+jhs/CxI8bc= +go.opentelemetry.io/proto/otlp v1.11.0 h1:5rrYs0Ykyj50sdU/JU0x8etU+LubXWb+gED6TbEdMIk= +go.opentelemetry.io/proto/otlp v1.11.0/go.mod h1:SmVizdCOAm3XBtG1g1NnOdhW6jtddT72hLMhv8VwA8E= +golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= +golang.org/x/arch v0.8.0 h1:3wRIsP3pM4yUptoR96otTUOXI367OS0+c9eeRi9doIc= +golang.org/x/arch v0.8.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= +golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= +golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= +golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d h1:FarXi840EJWSHYTN3ERkADbPWjl307+FGrA22KAVjjc= +google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d/go.mod h1:K/+WGbmBY7aNW1HDw1fJnKYo10i0DkAX6pows00dLig= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d h1:IL4hdHzcUv2l/gcg98/Rj3FbtE6axwqslOW8SW0C+S0= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.83.0 h1:JeNZEKJFbQxArAMl+hiytHauacDNqJUllNfmIMmpqnQ= +google.golang.org/grpc v1.83.0/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q= +gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA= +nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= +rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= diff --git a/internal/agent/agent.go b/internal/agent/agent.go new file mode 100644 index 0000000..c43b9fe --- /dev/null +++ b/internal/agent/agent.go @@ -0,0 +1,684 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "log" + "strings" + "sync" + "time" + "unicode" +) + +type Agent struct { + cfg Config + llm *LLMClient + store *SessionStore + tools *Toolset + docker *DockerConfigStore + apiCfg *APIConfigStore + coop *CoopManager + locks map[string]*sync.Mutex + locksMu sync.Mutex + liveMu sync.Mutex + live map[string]*liveReg +} + +// liveReg 记录某会话当前"在线"的 SSE 连接,用于推送后台任务完成通知。 +type liveReg struct { + emit func(Event) +} + +type Event struct { + Type string `json:"type"` + Data any `json:"data,omitempty"` +} + +func NewAgent(cfg Config, store *SessionStore, dockerCfg *DockerConfigStore, apiCfg *APIConfigStore) *Agent { + agent := &Agent{ + cfg: cfg, + llm: NewLLMClient(cfg, apiCfg), + store: store, + docker: dockerCfg, + apiCfg: apiCfg, + locks: make(map[string]*sync.Mutex), + live: make(map[string]*liveReg), + } + agent.tools = NewToolset(cfg, dockerCfg, apiCfg) + agent.coop = NewCoopManager() + agent.tools.coop = agent.coop + agent.coop.SetNotify(agent.notifyCoopDone) + return agent +} + +func (a *Agent) Run(ctx context.Context, sessionID, userContent string, emit func(Event)) error { + if strings.TrimSpace(userContent) == "" { + emit(Event{Type: "error", Data: map[string]any{"message": "消息内容不能为空"}}) + return errors.New("消息内容不能为空") + } + + lock := a.sessionLock(sessionID) + lock.Lock() + defer lock.Unlock() + + session, ok := a.store.Get(sessionID) + if !ok { + emit(Event{Type: "error", Data: map[string]any{"message": "会话不存在"}}) + return errors.New("会话不存在") + } + // 在私有副本上操作,避免修改持久化对象与其他会话的序列化产生数据竞争 + session = cloneSession(session) + + if session.Title == "" || session.Title == "新对话" { + session.Title = firstRunes(userContent, 30) + } + session.Messages = append(session.Messages, Message{ + Role: "user", + Content: stringPointer(userContent), + }) + a.store.Save(session) + + emit(Event{Type: "meta", Data: map[string]any{ + "session_id": session.ID, + "title": session.Title, + }}) + + history := session.Messages + compressed := false + iterations := 0 + emptyRetries := 0 + + for { + iterations++ + systemMessage := Message{ + Role: "system", + Content: stringPointer(a.systemPrompt(sessionID)), + } + apiMessages := append([]Message{systemMessage}, history...) + + if estimateMessages(apiMessages) > a.cfg.MaxContextTokens { + if summary, err := a.compressHistory(ctx, sessionID, history); err != nil { + history = truncateHistory(history, a.cfg.KeepRecentMessages) + compressed = true + emit(Event{Type: "notice", Data: map[string]any{ + "text": "上下文较长,自动压缩失败,已降级为截断早期对话: " + err.Error(), + }}) + } else { + history = summary + compressed = true + emit(Event{Type: "notice", Data: map[string]any{ + "text": "上下文较长,已自动压缩早期对话(分块摘要),关键信息会保留。", + }}) + } + } + + apiMessages = append([]Message{systemMessage}, history...) + assistant, err := a.streamAssistantReply(ctx, apiMessages, a.tools.Definitions(), emit) + if err != nil { + // 模型返回空内容:强制压缩上下文后重试,而非直接中断会话。 + // 长会话(如跑分循环累积大量工具输出)时模型容易因输入过大而返回空, + // 压缩后通常可恢复;多次仍空则优雅退出,保留会话状态供用户继续。 + if errors.Is(err, errEmptyResponse) { + emptyRetries++ + if emptyRetries <= 2 { + emit(Event{Type: "notice", Data: map[string]any{ + "text": "模型返回空内容,正在压缩上下文后重试…", + }}) + if summary, serr := a.compressHistory(ctx, sessionID, history); serr == nil { + history = summary + } else { + history = truncateHistory(history, a.cfg.KeepRecentMessages) + } + session.Messages = history + a.store.Save(session) + continue + } + session.Messages = history + a.store.Save(session) + emit(Event{Type: "notice", Data: map[string]any{ + "text": "模型多次返回空内容,本轮已停止。你可以继续发消息让我接着处理(程序不会退出)。", + }}) + emit(Event{Type: "done", Data: map[string]any{ + "compressed": compressed, + "iterations": iterations, + }}) + return nil + } + session.Messages = history + a.store.Save(session) + if errors.Is(err, context.Canceled) { + // 用户主动停止(前端 stop 按钮 / Esc):是正常中断,不按错误展示 + emit(Event{Type: "notice", Data: map[string]any{ + "text": "已停止本次生成,已完成的对话进度已保存,可继续发送新消息。", + }}) + } else { + emit(Event{Type: "error", Data: map[string]any{"message": err.Error()}}) + } + return err + } + emptyRetries = 0 + + history = append(history, assistant) + if len(assistant.ToolCalls) == 0 { + break + } + + // 本轮是否成功下发了协作任务:分发完成后本轮直接收尾, + // 避免控制层在 run_coop 之后继续空转(sleep 等待 / 反复轮询), + // coop 完成后会以 user 身份推送完成通知并启动新一轮处理。 + dispatched := false + for _, call := range assistant.ToolCalls { + emit(Event{Type: "tool", Data: map[string]any{ + "call_id": call.ID, + "name": call.Function.Name, + "input": call.Function.Arguments, + }}) + + result := a.tools.Execute(ctx, sessionID, call.Function.Name, call.Function.Arguments) + result.Output = truncateRunes(result.Output, a.cfg.MaxToolResultChars) + history = append(history, Message{ + Role: "tool", + ToolCallID: call.ID, + Content: stringPointer(result.Output), + }) + if call.Function.Name == "run_coop" && result.Success { + dispatched = true + } + + emit(Event{Type: "tool_result", Data: map[string]any{ + "call_id": call.ID, + "name": call.Function.Name, + "output": result.Output, + "success": result.Success, + }}) + } + + // 每轮迭代后落盘:长任务(如跑分循环)期间用户刷新页面也能看到 + // 已完成轮次的消息与工具结果,而不是只能等整轮 Run 结束。 + session.Messages = history + a.store.Save(session) + + if dispatched { + emit(Event{Type: "notice", Data: map[string]any{ + "text": "协作任务已全部下发,本轮收尾结束。任务完成后会自动以用户身份推送完成通知并启动新一轮处理,无需在等待上消耗本轮回合。", + }}) + break + } + + if iterations >= a.cfg.MaxIterations { + emit(Event{Type: "notice", Data: map[string]any{ + "text": fmt.Sprintf("已达到 %d 次工具调用上限,本轮先基于已有结果收尾。你可以继续发消息让我接着处理(程序不会退出)。", a.cfg.MaxIterations), + }}) + // 达到上限并不中断程序:用一次不带工具的收尾回答结束本回合, + // 避免在工具链中途戛然而止、用户得不到任何总结。 + closingSystem := Message{ + Role: "system", + Content: stringPointer(a.systemPrompt(sessionID) + + "\n\n注意:本轮已达到工具调用次数上限。现在必须直接输出最终回答:总结目前已确定的结果、未完成事项与建议的下一步,不要再调用任何工具。"), + } + if closing, err := a.streamAssistantReply(ctx, append([]Message{closingSystem}, history...), nil, emit); err == nil { + history = append(history, closing) + } + break + } + } + + session.Messages = history + a.store.Save(session) + emit(Event{Type: "done", Data: map[string]any{ + "compressed": compressed, + "iterations": iterations, + }}) + return nil +} + +// errEmptyResponse 表示模型流式返回既无文本也无工具调用。 +// 区别于普通错误:调用方应压缩上下文后重试,而非直接中断会话。 +var errEmptyResponse = errors.New("模型没有返回任何内容") + +func (a *Agent) streamAssistantReply(ctx context.Context, messages []Message, tools []ToolDefinition, emit func(Event)) (Message, error) { + const maxRetries = 2 + for attempt := 0; attempt <= maxRetries; attempt++ { + ch, err := a.llm.ChatStream(ctx, messages, tools) + if err != nil { + return Message{}, err + } + + var content strings.Builder + var calls []ToolCall + for chunk := range ch { + if chunk.Error != nil { + return Message{}, chunk.Error + } + if chunk.Content != "" { + content.WriteString(chunk.Content) + emit(Event{Type: "message", Data: map[string]any{"delta": chunk.Content}}) + } + for _, delta := range chunk.ToolCalls { + calls = mergeToolCall(calls, delta) + } + } + + for i := range calls { + if calls[i].ID == "" { + calls[i].ID = "call_" + newID() + } + } + + assistant := Message{ + Role: "assistant", + Content: stringPointer(content.String()), + } + if len(calls) > 0 { + assistant.ToolCalls = calls + if content.Len() == 0 { + assistant.Content = nil + } + } + if content.Len() == 0 && len(calls) == 0 { + if attempt < maxRetries { + emit(Event{Type: "notice", Data: map[string]any{ + "text": fmt.Sprintf("模型返回空内容,正在重试(第 %d/%d 次)…", attempt+1, maxRetries), + }}) + continue + } + return Message{}, errEmptyResponse + } + return assistant, nil + } + return Message{}, errEmptyResponse +} + +// compressHistory 压缩 history 中除最近 keep 条之外的全部早期消息。 +// 早期消息可能远超单次模型输入窗口,因此按压缩预算分块、逐块摘要后再合并; +// LLM 压缩失败时用 fallbackSummary 保留最近消息要点,避免直接整体截断。 +// 合并后若仍超窗口,对合并结果再压缩一层(递归收敛),极端情况才截断。 +func (a *Agent) compressHistory(ctx context.Context, sessionID string, history []Message) ([]Message, error) { + keep := a.cfg.KeepRecentMessages + if keep <= 0 { + keep = 12 + } + if len(history) <= keep { + return history, nil + } + old := history[:len(history)-keep] + recent := history[len(history)-keep:] + + summary, err := a.summarizeBlocks(ctx, old) + if err != nil { + // LLM 压缩失败:降级为文本要点摘要(纯文本,不产生孤儿 tool 消息) + summary = fallbackSummary(old) + } + merged := append([]Message{{ + Role: "system", + Content: stringPointer("[早期对话摘要]\n" + summary), + }}, recent...) + + // 压缩后校验:摘要 + 最近消息仍超窗口时,对合并结果再压缩一层(递归收敛) + if estimateMessages(merged) > a.cfg.MaxContextTokens { + if inner, inerr := a.compressHistory(ctx, sessionID, merged); inerr == nil { + return inner, nil + } + merged = truncateHistory(merged, keep) + } + return merged, nil +} + +// summarizeBlocks 把早期消息按单次压缩输入预算切成若干块,逐块交给 LLM 摘要, +// 最后把各块摘要按顺序合并返回。任一块失败即返回错误(由调用方降级)。 +func (a *Agent) summarizeBlocks(ctx context.Context, history []Message) (string, error) { + budget := a.compactionInputBudget() + var summaries []string + block := make([]Message, 0, 32) + blockTokens := 0 + flush := func() error { + if len(block) == 0 { + return nil + } + sum, cerr := a.llm.Compress(ctx, block) + block = block[:0] + blockTokens = 0 + if cerr != nil { + return cerr + } + summaries = append(summaries, sum) + return nil + } + for _, msg := range history { + tokens := estimateMessageTokens(msg) + if len(block) > 0 && blockTokens+tokens > budget { + if err := flush(); err != nil { + return "", err + } + } + block = append(block, msg) + blockTokens += tokens + } + if err := flush(); err != nil { + return "", err + } + return strings.Join(summaries, "\n\n"), nil +} + +// compactionInputBudget 返回单次 LLM 压缩请求的输入 token 预算: +// 总窗口减去压缩输出(CompactionTokens)与常规输出(StreamOutputTokens)的预留, +// 再留出安全余量,避免把输入塞满窗口导致请求被拒。 +func (a *Agent) compactionInputBudget() int { + window := a.cfg.MaxContextTokens + output := a.cfg.CompactionTokens + a.cfg.StreamOutputTokens + budget := window - output - window/8 + if budget < 4096 { + budget = 4096 + } + return budget +} + +// fallbackSummary 当 LLM 压缩失败时的兜底摘要:按时间顺序保留早期消息中 +// 最近若干条的用户目标/助手结论/工具结果要点,输出纯文本(不产生孤立的 +// tool 消息),保证后续请求协议合法、信息尽量不丢。 +func fallbackSummary(history []Message) string { + const maxLines = 10 + start := len(history) - maxLines + if start < 0 { + start = 0 + } + labels := map[string]string{ + "user": "用户", + "assistant": "助手", + "tool": "工具结果", + } + var lines []string + for i := start; i < len(history); i++ { + msg := history[i] + label, ok := labels[msg.Role] + if !ok { + label = msg.Role + } + text := "" + if msg.Content != nil { + text = strings.TrimSpace(*msg.Content) + } + if text == "" { + if len(msg.ToolCalls) > 0 { + text = "调用了 " + msg.ToolCalls[0].Function.Name + " 等工具" + } else { + continue + } + } + if countRunes(text) > 200 { + text = truncateRunes(text, 200) + } + lines = append(lines, fmt.Sprintf("[%s] %s", label, text)) + } + if len(lines) == 0 { + return "(早期对话内容较多,摘要失败后仅保留最近消息。)" + } + return "(早期对话压缩失败,以下为最近消息要点,更早内容已截断)\n" + strings.Join(lines, "\n") +} + +// truncateHistory 在模型压缩失败时兜底:仅保留最近消息,并丢弃开头孤立的 +// tool 结果消息,避免出现没有对应 assistant 调用声明的 tool 消息导致协议错乱。 +func truncateHistory(history []Message, keep int) []Message { + if len(history) <= keep { + return history + } + kept := append([]Message(nil), history[len(history)-keep:]...) + for len(kept) > 0 && kept[0].Role == "tool" { + kept = kept[1:] + } + return kept +} + +func (a *Agent) systemPrompt(sessionID string) string { + return fmt.Sprintf(`你是 BlackBean,运行在用户本机上的 AI Agent,正在协助用户完成真实任务。 +当前日期:%s +本会话工作目录(与其他会话隔离,相对路径与命令默认在此执行):%s +项目根目录(可读取,含 agent 自身文件、协作黑板产物等共享内容):%s + +工作方式: +1. 先理解用户目标,必要时先用工具查看文件、目录和运行结果,再给出准确答案。 +2. 你可以运行 Bash、读写文件、运行 Python。所有文件和命令默认都在本会话工作目录内进行,不要访问项目根目录之外的文件。 +3. 使用工具时必须根据真实输出继续推理,绝不能编造命令结果、文件内容或错误信息。 +4. 回答使用简洁的中文,重要代码、命令和路径用 Markdown 代码块呈现。 +5. 需要执行可能造成不可逆影响的操作前,先说明风险和影响,再谨慎执行。 +6. 如果一次工具调用没有解决问题,可以多次调用工具排查;不要在没有依据时下结论。 + + +【任务调度规则】 +- 解题类任务(CTF 题目、渗透测试、漏洞挖掘、Web 攻防、逆向、密码学、取证分析等)一律交由 coop 执行,你不负责具体解题,只做调度。 +- 你的调度职责: + 1. 从用户表述中提取足够信息:目标地址/文件、约束条件、平台提交规则(如 TSec Benchmark 的 submit API、BENCHMARK_TOKEN、unique_code 等); + 2. 调用 run_coop 工具,把任务描述与提交规则写清楚,交给 pi 协作单 Agent 执行; + 3. 任务完成后按完成通知中的结构化结果向用户汇报,并按约定规则提交 flag、回报得分。 +- 主动并发调度(默认行为,无需用户提醒):解题跑分类任务要最大化并发,始终把在跑的靶机/协作任务维持到平台并发上限(TSec Benchmark 通常为 3 个并发容器),不要一次只开一个靶机慢慢等: + 1. 开始解题时:先查看题目列表,选定多道未完成题目,用 start 接口一次性启动多个靶机,并逐一调用 run_coop 下发对应协作任务,把并发拉满; + 2. 每有任务完成/失败释放名额时:立即主动 close 已结束题目 → start 下一道未完成题目 → run_coop 下发,始终补满并发,不要空出名额等待; + 3. 只有平台已无可启动题目、或用户明确要求停止时才停止扩并发。 +- 主动检测状态(默认行为,无需用户提醒):在收到协作完成通知、下发新任务、或推进到关键节点时,主动用短命令(≤60s)查询平台进度(GET challenges)与运行中容器(docker_ps / run_bash),并向用户简要汇报「已通关题数/总分、运行中靶机、下一步计划」,不要等用户说「检测状态」才去查。 +- 调用 run_coop 成功下发目标后,本轮会话立即收尾结束:直接输出一句简短总结即可,不要原地等待、不要 sleep、不要反复轮询黑板。coop 完成后系统会自动以「用户」身份在会话中追加完成通知并启动新一轮处理,届时你再读取黑板产物、提交 flag、规划下一题。 +- coop 拿到 flag 后,必须按任务中约定的平台提交规则(如 TSec Benchmark 的 submit API、BENCHMARK_TOKEN、unique_code)用 curl 提交 flag,并把提交响应与得分回报给用户;提交失败要重试并说明原因,不能只汇报不提交。 +- 不要用 sleep 或超时超过 60 秒的长等待命令原地等待 coop 任务完成:coop 是异步运行的,任务完成后系统会自动推送「协作任务完成通知」并触发你继续处理,无需 sleep 阻塞。等待期间可以用短命令(≤60s)检查黑板进度,或直接推进其他不依赖该任务的工作(如规划下一题)。 +- 非解题类任务(文档撰写、代码开发、日常问答等)按常规方式由你自己完成,不必交给 coop。 +`, + time.Now().Format("2006-01-02"), a.tools.workspaceFor(sessionID), a.cfg.Workspace) +} + +func (a *Agent) sessionLock(sessionID string) *sync.Mutex { + a.locksMu.Lock() + defer a.locksMu.Unlock() + lock := a.locks[sessionID] + if lock == nil { + lock = &sync.Mutex{} + a.locks[sessionID] = lock + } + return lock +} + +// ForgetSession 释放会话持有的运行时资源(会话锁与 live 注册)。 +// 在删除会话时调用,避免 locks / live map 随会话累积只增不减。 +func (a *Agent) ForgetSession(sessionID string) { + a.locksMu.Lock() + delete(a.locks, sessionID) + a.locksMu.Unlock() + + a.liveMu.Lock() + delete(a.live, sessionID) + a.liveMu.Unlock() +} + +// RegisterLive 把某个 SSE 连接注册为会话的"在线接收者",供后台协作任务完成 +func (a *Agent) RegisterLive(sessionID string, emit func(Event)) func() { + a.liveMu.Lock() + reg := &liveReg{emit: emit} + a.live[sessionID] = reg + a.liveMu.Unlock() + return func() { + a.liveMu.Lock() + if a.live[sessionID] == reg { + delete(a.live, sessionID) + } + a.liveMu.Unlock() + } +} + +// CoopTasks 返回协作任务列表;sessionID 非空时只返回该会话的任务, +// 供 Web 页面实时展示容器运行状态。 +func (a *Agent) CoopTasks(sessionID string) []*CoopTask { + return a.coop.List(sessionID) +} + +// liveEmitter 返回动态事件发射器:每次发送时实时查询该会话当前在线的连接。 +// 与启动时一次性捕获 emit 相比,即使汇报轮在页面连接之前启动,页面打开后 +// 也能实时收到流式事件(修复"汇报轮事件前端收不到"的问题)。 +func (a *Agent) liveEmitter(sessionID string) func(Event) { + return func(event Event) { + a.liveMu.Lock() + reg := a.live[sessionID] + a.liveMu.Unlock() + if reg != nil { + reg.emit(event) + } + } +} + +// notifyCoopDone 是后台协作任务完成后的回调:给会话注入一条 +// "【协作任务完成通知】"用户消息,并触发主 agent 新一轮处理 +// (有在线 SSE 连接则流式推送,否则在后台静默完成并持久化结果)。 +// 通知基于 supervisor 的结构化结果(CoopResult)渲染,成功时携带 flag/摘要/ +// 产物清单,失败时携带明确的续跑指令(关闭靶机 → 开下一题 → 重新下发 coop)。 +// 汇报轮失败会在会话中持久化错误信息,避免"通知已到但流程静默停摆"。 +func (a *Agent) notifyCoopDone(task *CoopTask) { + notif := buildCoopNotification(task) + + emit := a.liveEmitter(task.SessionID) + // 汇报轮加超时兜底:后台自动触发的 Run 不能永久持锁阻塞后续所有通知 + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute) + emit(Event{Type: "turn_start", Data: map[string]any{"user": notif}}) + // 使用独立后台上下文:即使前端已关闭页面,汇报轮仍会完成并持久化。 + // 注意:cancel 必须在 Run 结束后调用(defer 到 goroutine 内), + // 否则本函数一返回 ctx 即被取消,汇报轮的首个 LLM 请求会立即报 + // "Post .../v1/messages: context canceled"(历史 bug:会话停在完成通知处)。 + go func() { + defer cancel() + if err := a.Run(ctx, task.SessionID, notif, emit); err != nil { + // 汇报轮失败兜底:把错误持久化到会话,避免静默断流 + if session, ok := a.store.Get(task.SessionID); ok { + fallback := "[系统] 协作完成通知的自动处理失败(" + err.Error() + + ")。请先读取黑板产物(" + task.Blackboard + ")向用户汇报,再按 TSec 流程继续推进:close 已通关题目释放名额 → start 下一题 → run_coop 下发协作。" + session.Messages = append(session.Messages, Message{Role: "user", Content: stringPointer(fallback)}) + a.store.Save(session) + } + log.Printf("[coop] 汇报轮失败 session=%s task=%s err=%v", task.SessionID, task.ID, err) + } + }() +} + +// buildCoopNotification 根据任务的结构化结果渲染完成通知。 +// 有 result.json 时按 status 判定成功/失败并附带 flag/产物清单; +// 缺失时回退到容器 stdout(老版本兼容)。 +func buildCoopNotification(task *CoopTask) string { + success := false + status := "完成" + var summary, flag string + var artifacts []string + + if task.Result != nil { + res := task.Result + switch res.Status { + case "solved": + success = true + status = "完成" + case "unsolved": + status = "未解出" + case "timeout": + status = "超时" + default: + status = "失败" + } + summary, flag, artifacts = res.Summary, res.Flag, res.Artifacts + } else { + // 老版本 / result.json 缺失时回退到容器日志 + detail := task.Output + if detail == "" { + detail = task.Error + } + failed := task.Error != "" || task.ExitCode != 0 || + strings.Contains(detail, "未解出") || strings.Contains(detail, "未完成") + if failed { + status = "未解出/失败" + } else { + success = true + } + summary = detail + } + + var b strings.Builder + fmt.Fprintf(&b, "【协作任务完成通知】你之前发起的 pi 协作任务(ID: %s)已%s。\n", task.ID, status) + if success && flag != "" { + fmt.Fprintf(&b, "\n【flag】%s\n", flag) + } + if strings.TrimSpace(summary) != "" { + b.WriteString("\n" + truncateRunes(strings.TrimSpace(summary), 1200) + "\n") + } + if task.Blackboard != "" { + fmt.Fprintf(&b, "\n黑板产物目录:%s", task.Blackboard) + if len(artifacts) > 0 { + b.WriteString("\n关键产物:") + for _, artifact := range artifacts { + if name := strings.TrimSpace(artifact); name != "" { + b.WriteString("\n- " + name) + } + } + } + b.WriteString("\n可选用 read_file / list_directory 查看黑板产物,向用户汇报结果与关键文件。") + } + if success { + b.WriteString("\n若产物中包含 flag 且尚未提交:请读取黑板下 task.md(内含平台提交规则 / BENCHMARK_TOKEN / unique_code),立即用 curl 向平台提交 flag 并给出提交响应与得分;不要只汇报而不提交。") + b.WriteString("\n随后按 TSec 流程持续推进以拿更高总分:close 已通关题目释放名额 → 主动补满并发(一次性 start 多道未完成题目并逐一 run_coop,把在跑靶机数拉满到平台上限)→ 不要空出名额等待;若所有题目已完成或平台任务超时,则停止并向用户汇报总分。") + } else { + b.WriteString("\n该任务未解出/失败,请按 TSec 标准流程继续推进,不要空等:") + b.WriteString("\n1) 用 POST {BENCHMARK_BASE_URL}/openapi/v1/challenges/close?unique_code=<该题编号> 关闭当前靶机容器释放名额(该题编号为 " + task.ChallengeCode + ",可核对黑板 task.md);") + b.WriteString("\n2) 用 GET {BENCHMARK_BASE_URL}/openapi/v1/challenges 查看剩余未完成题目,一次性选定多道未完成题目补满并发;") + b.WriteString("\n3) 用 POST {BENCHMARK_BASE_URL}/openapi/v1/challenges/start?unique_code=<新题> 逐个启动新靶机,把在跑靶机数拉满到平台并发上限;") + b.WriteString("\n4) 对每道新题调用 run_coop 工具下发协作解题,不要只开一个靶机空等。") + b.WriteString("\n若所有题目已完成或平台任务已结束(接口持续 invalid_state),则停止并向用户汇报总分。") + } + return b.String() +} + +func mergeToolCall(calls []ToolCall, delta ToolCallDelta) []ToolCall { + index := delta.Index + if index < 0 { + index = len(calls) + } + for len(calls) <= index { + calls = append(calls, ToolCall{Type: "function"}) + } + call := &calls[index] + if delta.ID != "" { + call.ID = delta.ID + } + if delta.Name != "" { + call.Function.Name = delta.Name + } + call.Function.Arguments += delta.ArgumentsDelta + return calls +} + +func estimateMessages(messages []Message) int { + total := 0 + for _, message := range messages { + total += estimateMessageTokens(message) + } + return total +} + +// estimateMessageTokens 估算单条消息的 token 数(含 tool_calls 与消息开销)。 +func estimateMessageTokens(message Message) int { + total := 0 + if message.Content != nil { + total += estimateTokens(*message.Content) + } + for _, call := range message.ToolCalls { + total += estimateTokens(call.Function.Name + call.Function.Arguments) + } + total += 8 + return total +} + +func estimateTokens(value string) int { + runes := []rune(value) + hanCount := 0 + for _, r := range runes { + if unicode.Is(unicode.Han, r) { + hanCount++ + } + } + other := len(runes) - hanCount + return int(float64(other)/4.0+float64(hanCount)*0.8) + 4 +} + +func stringPointer(value string) *string { + return &value +} diff --git a/internal/agent/api_config.go b/internal/agent/api_config.go new file mode 100644 index 0000000..56e7712 --- /dev/null +++ b/internal/agent/api_config.go @@ -0,0 +1,258 @@ +package agent + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "sync" +) + +// 支持的接口类型。 +const ( + ProviderOpenAI = "openai" + ProviderAnthropic = "anthropic" +) + +// 协作 worker 管理模式。 +const ( + // CoopModeDocker 通过 Docker 容器管理 worker(需本机/远程 Docker daemon)。 + CoopModeDocker = "docker" + // CoopModeLocal 通过本地子进程管理 worker(无需 Docker,托管沙箱等场景)。 + CoopModeLocal = "local" +) + +// LLM 配置环境变量名。环境变量优先级最高,覆盖 api-config.json 用户配置与内置默认值。 +// 适用于托管沙箱等场景:平台注入环境变量,Agent 启动即生效,无需通过 Web 设置页配置。 +const ( + EnvLLMAPIKey = "LLM_API_KEY" + EnvLLMBaseURL = "LLM_BASE_URL" + EnvLLMModel = "LLM_MODEL" + EnvLLMProvider = "LLM_PROVIDER" + EnvLLMEngine = "LLM_ENGINE" + EnvCoopMode = "COOP_MODE" +) + +// APIConfig 是用户可配置的 LLM API 连接信息。 +type APIConfig struct { + // APIKey 是模型服务商 API Key。 + APIKey string `json:"api_key,omitempty"` + // BaseURL 是模型接口地址。 + BaseURL string `json:"base_url,omitempty"` + // Model 是模型名称。 + Model string `json:"model,omitempty"` + // Provider 是接口类型:openai(OpenAI 兼容)或 anthropic(Anthropic Messages API)。 + Provider string `json:"provider,omitempty"` + // Engine 是默认协作引擎:pi(默认)/ pigo / claude。 + // 仅影响 run_coop 未显式指定 engine 时的默认值,见 tools_coop.go。 + Engine string `json:"engine,omitempty"` + // CoopMode 是协作 worker 管理模式:local(默认,本地子进程)/ docker(Docker 容器)。 + // local 模式无需 Docker,适用于托管沙箱等无 Docker 环境;docker 模式需本机/远程 Docker daemon。 + CoopMode string `json:"coop_mode,omitempty"` +} + +// APIConfigStore 持久化用户的 LLM API 配置,未配置的字段回退到默认值(环境变量)。 +type APIConfigStore struct { + mu sync.RWMutex + path string + config APIConfig + defaults APIConfig +} + +func NewAPIConfigStore(dataDir string, defaults APIConfig) (*APIConfigStore, error) { + store := &APIConfigStore{ + path: filepath.Join(dataDir, "api-config.json"), + defaults: defaults, + } + store.load() + return store, nil +} + +func (s *APIConfigStore) load() { + data, err := os.ReadFile(s.path) + if err != nil { + return + } + if len(data) == 0 { + return + } + var config APIConfig + if err := json.Unmarshal(data, &config); err != nil { + return + } + s.config = config +} + +func (s *APIConfigStore) persist() error { + data, err := json.MarshalIndent(s.config, "", " ") + if err != nil { + return err + } + tmp := s.path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return err + } + return os.Rename(tmp, s.path) +} + +// APIKey 返回生效的 API Key。 +// 优先级:环境变量 LLM_API_KEY > api-config.json > 内置默认值。 +func (s *APIConfigStore) APIKey() string { + if v := strings.TrimSpace(os.Getenv(EnvLLMAPIKey)); v != "" { + return v + } + s.mu.RLock() + defer s.mu.RUnlock() + if s.config.APIKey != "" { + return s.config.APIKey + } + return s.defaults.APIKey +} + +// BaseURL 返回生效的接口地址。 +// 优先级:环境变量 LLM_BASE_URL > api-config.json > 内置默认值。 +func (s *APIConfigStore) BaseURL() string { + if v := strings.TrimSpace(os.Getenv(EnvLLMBaseURL)); v != "" { + return v + } + s.mu.RLock() + defer s.mu.RUnlock() + if s.config.BaseURL != "" { + return s.config.BaseURL + } + return s.defaults.BaseURL +} + +// Model 返回生效的模型名称。 +// 优先级:环境变量 LLM_MODEL > api-config.json > 内置默认值。 +func (s *APIConfigStore) Model() string { + if v := strings.TrimSpace(os.Getenv(EnvLLMModel)); v != "" { + return v + } + s.mu.RLock() + defer s.mu.RUnlock() + if s.config.Model != "" { + return s.config.Model + } + return s.defaults.Model +} + +// IsAPIKeyConfigured 报告是否已配置 API Key(含环境变量、用户配置与内置默认值)。 +// LLM 客户端实际会按 APIKey() 的优先级(环境变量 > api-config.json > 默认值)取用, +// 因此只要最终能拿到非空 Key 就视为已配置,避免误判导致无法对话。 +func (s *APIConfigStore) IsAPIKeyConfigured() bool { + if v := strings.TrimSpace(os.Getenv(EnvLLMAPIKey)); v != "" { + return true + } + s.mu.RLock() + defer s.mu.RUnlock() + return s.config.APIKey != "" || s.defaults.APIKey != "" +} + +// Provider 返回生效的接口类型。 +// 优先级:环境变量 LLM_PROVIDER > api-config.json > 内置默认值 > openai。 +func (s *APIConfigStore) Provider() string { + if v := strings.TrimSpace(os.Getenv(EnvLLMProvider)); v != "" { + if v == ProviderOpenAI || v == ProviderAnthropic { + return v + } + } + s.mu.RLock() + defer s.mu.RUnlock() + if s.config.Provider != "" { + return s.config.Provider + } + if s.defaults.Provider != "" { + return s.defaults.Provider + } + return ProviderOpenAI +} + +// Engine 返回生效的默认协作引擎:pi(默认)/ pigo / claude。 +// 优先级:环境变量 LLM_ENGINE > api-config.json > 内置默认值 > pi。 +// 仅在 run_coop 未显式指定 engine 时使用。 +func (s *APIConfigStore) Engine() string { + if v := strings.TrimSpace(os.Getenv(EnvLLMEngine)); v != "" { + switch v { + case "pi", "pigo", "claude": + return v + } + } + s.mu.RLock() + defer s.mu.RUnlock() + switch s.config.Engine { + case "pi", "pigo", "claude": + return s.config.Engine + } + if s.defaults.Engine != "" { + return s.defaults.Engine + } + return "pi" +} + +// CoopMode 返回生效的协作 worker 管理模式:local(默认)/ docker。 +// 优先级:环境变量 COOP_MODE > api-config.json > 内置默认值 > local。 +// local 模式通过本地子进程运行 worker(无需 Docker),docker 模式通过 Docker 容器运行。 +func (s *APIConfigStore) CoopMode() string { + if v := strings.TrimSpace(os.Getenv(EnvCoopMode)); v != "" { + switch v { + case CoopModeDocker, CoopModeLocal: + return v + } + } + s.mu.RLock() + defer s.mu.RUnlock() + switch s.config.CoopMode { + case CoopModeDocker, CoopModeLocal: + return s.config.CoopMode + } + if s.defaults.CoopMode != "" { + return s.defaults.CoopMode + } + return CoopModeLocal +} + +// Update 更新用户配置。每个参数为 nil 表示不修改该字段; +// 非 nil(含空字符串)表示设置或清除该字段(空串=清除,回退默认)。 +func (s *APIConfigStore) Update(apiKey, baseURL, model, provider, engine, coopMode *string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if apiKey != nil { + s.config.APIKey = *apiKey + } + if baseURL != nil { + s.config.BaseURL = *baseURL + } + if model != nil { + s.config.Model = *model + } + if provider != nil { + value := strings.TrimSpace(*provider) + if value != ProviderOpenAI && value != ProviderAnthropic { + value = "" + } + s.config.Provider = value + } + if engine != nil { + value := strings.TrimSpace(*engine) + switch value { + case "pi", "pigo", "claude": + // 合法值 + default: + value = "" + } + s.config.Engine = value + } + if coopMode != nil { + value := strings.TrimSpace(*coopMode) + switch value { + case CoopModeDocker, CoopModeLocal: + // 合法值 + default: + value = "" + } + s.config.CoopMode = value + } + return s.persist() +} diff --git a/internal/agent/compress_test.go b/internal/agent/compress_test.go new file mode 100644 index 0000000..7ad09ee --- /dev/null +++ b/internal/agent/compress_test.go @@ -0,0 +1,71 @@ +package agent + +import ( + "strings" + "testing" +) + +// TestFallbackSummary 验证 LLM 压缩失败时的兜底摘要: +// 纯文本、保留最近消息要点、无孤立 tool 消息。 +func TestFallbackSummary(t *testing.T) { + history := []Message{ + {Role: "user", Content: stringPointer("第一轮任务:分析日志")}, + {Role: "assistant", Content: stringPointer("已确认方案,开始执行")}, + {Role: "assistant", ToolCalls: []ToolCall{{ID: "c1", Type: "function", Function: ToolFunctionCall{Name: "run_bash", Arguments: `{"command":"ls"}`}}}}, + {Role: "tool", ToolCallID: "c1", Content: stringPointer("file.txt")}, + {Role: "user", Content: stringPointer("继续,报告结果")}, + } + summary := fallbackSummary(history) + if summary == "" { + t.Fatal("fallbackSummary 不应返回空") + } + if strings.Contains(summary, "tool_calls") || strings.Contains(summary, "\"id\"") { + t.Fatalf("fallbackSummary 必须为纯文本,不得包含结构化 tool_calls: %s", summary) + } + if !strings.Contains(summary, "用户") || !strings.Contains(summary, "助手") || !strings.Contains(summary, "工具结果") { + t.Fatalf("fallbackSummary 应包含用户/助手/工具结果要点: %s", summary) + } + if !strings.Contains(summary, "调用了 run_bash") { + t.Fatalf("fallbackSummary 应记录工具调用名: %s", summary) + } +} + +// TestCompactionInputBudget 验证单次压缩输入预算随窗口放大而放大, +// 且始终留出输出与安全余量。 +func TestCompactionInputBudget(t *testing.T) { + a := &Agent{cfg: Config{MaxContextTokens: 96000, CompactionTokens: 4096, StreamOutputTokens: 4096}} + budget := a.compactionInputBudget() + if budget <= 0 || budget >= 96000 { + t.Fatalf("预算应介于 (0, 96000),got %d", budget) + } + // 窗口放大后预算也应放大 + a2 := &Agent{cfg: Config{MaxContextTokens: 48000, CompactionTokens: 4096, StreamOutputTokens: 4096}} + if a2.compactionInputBudget() >= budget { + t.Fatalf("窗口更大的预算应更大,got %d vs %d", a2.compactionInputBudget(), budget) + } + // 极端小窗口时兜底下限 4096 + a3 := &Agent{cfg: Config{MaxContextTokens: 2000, CompactionTokens: 512, StreamOutputTokens: 512}} + if got := a3.compactionInputBudget(); got < 4096 { + t.Fatalf("小窗口预算应有下限,got %d", got) + } +} + +// TestEstimateMessageTokens 验证单条消息估算包含 content、tool_calls 与开销。 +func TestEstimateMessageTokens(t *testing.T) { + m := Message{ + Role: "assistant", + Content: stringPointer("你好"), + ToolCalls: []ToolCall{{ + ID: "c1", Type: "function", + Function: ToolFunctionCall{Name: "read_file", Arguments: `{"path":"a.txt"}`}, + }}, + } + n := estimateMessageTokens(m) + if n <= 0 { + t.Fatalf("估算应大于 0,got %d", n) + } + plain := estimateMessageTokens(Message{Role: "user", Content: stringPointer("你好")}) + if n <= plain { + t.Fatalf("含 tool_calls 的消息估算应大于纯文本,got %d vs %d", n, plain) + } +} diff --git a/internal/agent/config.go b/internal/agent/config.go new file mode 100644 index 0000000..b6109a8 --- /dev/null +++ b/internal/agent/config.go @@ -0,0 +1,137 @@ +package agent + +import ( + "bufio" + "os" + "path/filepath" + "strconv" + "strings" + "time" +) + +type Config struct { + APIKey string + BaseURL string + Model string + Port string + Workspace string + DataDir string + // AuthCode 是 Web/API/WS 访问授权码。为空(默认)表示不启用授权; + // 非空(Docker 部署时通过环境变量 AGENT_AUTH_CODE 注入)表示启用: + // 访问者必须输入正确授权码才能使用本实例。 + AuthCode string + MaxContextTokens int + KeepRecentMessages int + MaxToolResultChars int + MaxIterations int + ToolTimeout time.Duration + RequestTimeout time.Duration + StreamOutputTokens int + CompactionTokens int + Temperature float64 +} + +func LoadConfig() Config { + loadDotEnv(".env") + + workspace := getenv("AGENT_WORKSPACE", "") + if workspace == "" { + workspace, _ = os.Getwd() + } + absWorkspace, err := filepath.Abs(workspace) + if err != nil { + absWorkspace = workspace + } + dataDir := getenv("AGENT_DATA_DIR", filepath.Join(absWorkspace, "data")) + if !filepath.IsAbs(dataDir) { + dataDir = filepath.Join(absWorkspace, dataDir) + } + + cfg := Config{ + APIKey: os.Getenv("SILICONFLOW_API_KEY"), + BaseURL: getenv("SILICONFLOW_BASE_URL", "https://api.siliconflow.cn/v1/chat/completions"), + Model: getenv("AGENT_MODEL", "Qwen/Qwen3-32B"), + Port: getenv("AGENT_PORT", "8080"), + Workspace: absWorkspace, + DataDir: dataDir, + // 全面放开限制以适配 DeepSeek(128K 上下文、大输出)跑复杂任务: + // - MaxContextTokens 给到 1M,确保 compactionInputBudget 公式(window - Compaction - Stream - window/8)仍有充足余量 + // - StreamOutputTokens 提升到 65536:直接作为 max_tokens 传给 LLM(见 llm.go),是放开主 Agent 输出能力的关键 + // - CompactionTokens 提升到 16384:压缩摘要允许足够长度,避免历史信息丢失 + // - MaxToolResultChars 提升到 100000:复杂任务的脚本/日志能完整传回 LLM,不被截断 + // - KeepRecentMessages 提升到 50:压缩后保留更多历史,复杂多步任务上下文不丢 + // - ToolTimeout/RequestTimeout 提升到 600s:编译/长命令/深度推理有充足时间 + MaxContextTokens: getenvInt("AGENT_MAX_CONTEXT_TOKENS", 1000000), + KeepRecentMessages: getenvInt("AGENT_KEEP_RECENT_MESSAGES", 50), + MaxToolResultChars: getenvInt("AGENT_MAX_TOOL_RESULT_CHARS", 100000), + MaxIterations: getenvInt("AGENT_MAX_ITERATIONS", 50), + ToolTimeout: time.Duration(getenvInt("AGENT_TOOL_TIMEOUT_SECONDS", 600)) * time.Second, + RequestTimeout: time.Duration(getenvInt("AGENT_REQUEST_TIMEOUT_SECONDS", 600)) * time.Second, + StreamOutputTokens: getenvInt("AGENT_STREAM_OUTPUT_TOKENS", 65536), + CompactionTokens: getenvInt("AGENT_COMPACTION_TOKENS", 16384), + Temperature: getenvFloat("AGENT_TEMPERATURE", 0.3), + // 授权码默认关闭:仅当 Docker 部署时显式注入 AGENT_AUTH_CODE 才启用访问授权。 + AuthCode: getenv("AGENT_AUTH_CODE", ""), + } + return cfg +} + +func loadDotEnv(path string) { + file, err := os.Open(path) + if err != nil { + return + } + defer file.Close() + + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + key, value, ok := strings.Cut(line, "=") + if !ok { + continue + } + key = strings.TrimSpace(key) + value = strings.TrimSpace(value) + value = strings.Trim(value, `"'`) + if key == "" { + continue + } + if _, exists := os.LookupEnv(key); !exists { + _ = os.Setenv(key, value) + } + } +} + +func getenv(key, fallback string) string { + if value := os.Getenv(key); value != "" { + return value + } + return fallback +} + +func getenvInt(key string, fallback int) int { + value := os.Getenv(key) + if value == "" { + return fallback + } + parsed, err := strconv.Atoi(value) + if err != nil { + return fallback + } + return parsed +} + +func getenvFloat(key string, fallback float64) float64 { + value := os.Getenv(key) + if value == "" { + return fallback + } + parsed, err := strconv.ParseFloat(value, 64) + if err != nil { + return fallback + } + return parsed +} diff --git a/internal/agent/coop_async.go b/internal/agent/coop_async.go new file mode 100644 index 0000000..db2019a --- /dev/null +++ b/internal/agent/coop_async.go @@ -0,0 +1,180 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "sort" + "sync" + "time" + + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/filters" + "github.com/docker/docker/client" +) + +// CoopTask 记录一次后台运行的 pi 协作任务。 +// 任务由 run_coop 工具异步启动,完成后通过 CoopManager.Complete 通知主 agent。 +type CoopTask struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + ContainerID string `json:"container_id,omitempty"` + // ProcessID 是本地子进程模式的 worker 进程 PID(docker 模式为 0)。 + ProcessID int `json:"process_id,omitempty"` + Status string `json:"status"` // running | done + ExitCode int `json:"exit_code,omitempty"` + Output string `json:"output,omitempty"` + Blackboard string `json:"blackboard,omitempty"` + Error string `json:"error,omitempty"` + RoundMax int `json:"round_max"` + CreatedAt time.Time `json:"created_at"` + FinishedAt time.Time `json:"finished_at,omitempty"` + // Result 是 supervisor 写入黑板 result.json 的结构化结果(存在时优先于 Output)。 + Result *CoopResult `json:"result,omitempty"` + // ChallengeCode 任务描述的题目编号(如 c-06),用于通知控制层关闭/切换靶机 + ChallengeCode string `json:"challenge_code,omitempty"` +} + +// CoopResult 是 pi 协作容器 supervisor 生成的标准结果($BB/result.json), +// 完成通知与主 agent 汇报轮直接使用这些结构化字段,不再依赖截断的容器 stdout。 +type CoopResult struct { + Status string `json:"status"` // solved | unsolved | error | timeout + ExitCode int `json:"exit_code"` + Summary string `json:"summary"` + Flag string `json:"flag"` + Artifacts []string `json:"artifacts"` +} + +// Solved 返回任务是否成功解出/交付完成。 +func (r *CoopResult) Solved() bool { + return r != nil && r.Status == "solved" +} + +// CoopManager 管理后台协作任务的生命周期,并在任务完成时回调 Agent, +// 由 Agent 决定是否唤醒主 agent 汇报结果。 +type CoopManager struct { + mu sync.Mutex + tasks map[string]*CoopTask + notify func(task *CoopTask) +} + +func NewCoopManager() *CoopManager { + return &CoopManager{tasks: make(map[string]*CoopTask)} +} + +// SetNotify 注册任务完成回调。 +func (m *CoopManager) SetNotify(fn func(task *CoopTask)) { + m.mu.Lock() + m.notify = fn + m.mu.Unlock() +} + +// Register 登记一个正在运行的后台任务。 +func (m *CoopManager) Register(task *CoopTask) { + m.mu.Lock() + m.tasks[task.ID] = task + m.mu.Unlock() +} + +// Get 按任务 ID 查询任务。 +func (m *CoopManager) Get(id string) (*CoopTask, bool) { + m.mu.Lock() + defer m.mu.Unlock() + task, ok := m.tasks[id] + return task, ok +} + +// List 返回全部任务,按创建时间倒序;sessionID 非空时只返回该会话的任务。 +func (m *CoopManager) List(sessionID string) []*CoopTask { + m.mu.Lock() + defer m.mu.Unlock() + out := make([]*CoopTask, 0, len(m.tasks)) + for _, task := range m.tasks { + if sessionID == "" || task.SessionID == sessionID { + out = append(out, task) + } + } + sort.Slice(out, func(i, j int) bool { + return out[i].CreatedAt.After(out[j].CreatedAt) + }) + return out +} + +// Complete 标记任务完成并触发完成回调(回调在锁外执行,避免死锁)。 +func (m *CoopManager) Complete(task *CoopTask) { + m.mu.Lock() + task.Status = "done" + task.FinishedAt = time.Now() + m.tasks[task.ID] = task + notify := m.notify + m.mu.Unlock() + if notify != nil { + notify(task) + } +} + +// CurrentCoopRound 统计黑板 logs 目录中已出现的最大协作轮次。 +// 单 agent 版日志命名为 round-N.log;兼容旧版 a-round-N.log / b-round-N.log。 +func CurrentCoopRound(blackboardDir string) int { + if blackboardDir == "" { + return 0 + } + entries, err := os.ReadDir(filepath.Join(blackboardDir, "logs")) + if err != nil { + return 0 + } + maxRound := 0 + for _, entry := range entries { + var round int + if _, err := fmt.Sscanf(entry.Name(), "round-%d.log", &round); err == nil { + if round > maxRound { + maxRound = round + } + continue + } + if _, err := fmt.Sscanf(entry.Name(), "a-round-%d.log", &round); err == nil { + if round > maxRound { + maxRound = round + } + } + if _, err := fmt.Sscanf(entry.Name(), "b-round-%d.log", &round); err == nil { + if round > maxRound { + maxRound = round + } + } + } + return maxRound +} + +// ContainerInspectStatus 查询容器当前的运行状态与人类可读状态描述 +// (如 "running" + "Up 2 minutes" / "exited" + "Exited (0)"), +// 供 Web 页面实时展示协作容器状态。 +func ContainerInspectStatus(socket, containerID string) (state, status string, err error) { + if containerID == "" { + return "", "", errors.New("缺少容器 ID") + } + cli, err := client.NewClientWithOpts( + client.WithHost(socket), + client.WithAPIVersionNegotiation(), + ) + if err != nil { + return "", "", err + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + list, err := cli.ContainerList(ctx, container.ListOptions{ + All: true, + Filters: filters.NewArgs(filters.Arg("id", containerID)), + }) + if err != nil { + return "", "", err + } + if len(list) == 0 { + return "", "", errors.New("容器不存在") + } + return list[0].State, list[0].Status, nil +} diff --git a/internal/agent/coop_test.go b/internal/agent/coop_test.go new file mode 100644 index 0000000..bf5f99d --- /dev/null +++ b/internal/agent/coop_test.go @@ -0,0 +1,85 @@ +package agent + +import ( + "context" + "testing" + "time" +) + +func TestCoopBaseURL(t *testing.T) { + tests := []struct { + name string + provider string + engine string + baseURL string + want string + }{ + // ---- pi / claude(用 @anthropic-ai/sdk,自带 /v1/messages,不能补 /v1)---- + // Anthropic:剥离 /v1,避免双重 /v1/v1/messages + {"pi/anthropic bare", "anthropic", "pi", "https://api.deepseek.com/anthropic", "https://api.deepseek.com/anthropic"}, + {"pi/anthropic trailing slash", "anthropic", "pi", "https://api.deepseek.com/anthropic/", "https://api.deepseek.com/anthropic"}, + {"pi/anthropic with /v1", "anthropic", "pi", "https://api.deepseek.com/anthropic/v1", "https://api.deepseek.com/anthropic"}, + {"pi/anthropic with /v1/messages", "anthropic", "pi", "https://api.deepseek.com/anthropic/v1/messages", "https://api.deepseek.com/anthropic"}, + {"pi/anthropic with /messages", "anthropic", "pi", "https://api.deepseek.com/anthropic/messages", "https://api.deepseek.com/anthropic"}, + {"pi/anthropic anthropic.com", "anthropic", "pi", "https://api.anthropic.com", "https://api.anthropic.com"}, + // claude 行为与 pi 一致 + {"claude/anthropic bare", "anthropic", "claude", "https://api.deepseek.com/anthropic", "https://api.deepseek.com/anthropic"}, + {"claude/anthropic with /v1", "anthropic", "claude", "https://api.deepseek.com/anthropic/v1", "https://api.deepseek.com/anthropic"}, + + // ---- pigo(anthropicCompatDriver 只追加 /messages,需补 /v1)---- + {"pigo/anthropic bare", "anthropic", "pigo", "https://api.deepseek.com/anthropic", "https://api.deepseek.com/anthropic/v1"}, + {"pigo/anthropic trailing slash", "anthropic", "pigo", "https://api.deepseek.com/anthropic/", "https://api.deepseek.com/anthropic/v1"}, + {"pigo/anthropic with /v1", "anthropic", "pigo", "https://api.deepseek.com/anthropic/v1", "https://api.deepseek.com/anthropic/v1"}, + {"pigo/anthropic with /v1/messages", "anthropic", "pigo", "https://api.deepseek.com/anthropic/v1/messages", "https://api.deepseek.com/anthropic/v1"}, + {"pigo/anthropic with /messages", "anthropic", "pigo", "https://api.deepseek.com/anthropic/messages", "https://api.deepseek.com/anthropic/v1"}, + {"pigo/anthropic anthropic.com", "anthropic", "pigo", "https://api.anthropic.com", "https://api.anthropic.com/v1"}, + + // ---- OpenAI 协议(三种 engine 行为一致:保留 /v1,剥离 /chat/completions)---- + {"pi/openai full url", "openai", "pi", "https://api.siliconflow.cn/v1/chat/completions", "https://api.siliconflow.cn/v1"}, + {"pi/openai with /v1", "openai", "pi", "https://api.siliconflow.cn/v1", "https://api.siliconflow.cn/v1"}, + {"pi/openai bare", "openai", "pi", "https://api.siliconflow.cn", "https://api.siliconflow.cn"}, + {"pigo/openai full url", "openai", "pigo", "https://api.siliconflow.cn/v1/chat/completions", "https://api.siliconflow.cn/v1"}, + {"pigo/openai bare", "openai", "pigo", "https://api.siliconflow.cn", "https://api.siliconflow.cn"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := coopBaseURL(tt.provider, tt.baseURL, tt.engine) + if got != tt.want { + t.Errorf("coopBaseURL(%q, %q, %q) = %q, want %q", tt.provider, tt.baseURL, tt.engine, got, tt.want) + } + }) + } +} + +func TestRunCoopErrorPaths(t *testing.T) { + apiCfg := &APIConfigStore{} + apiCfg.config = APIConfig{APIKey: "test-key", BaseURL: "https://api.deepseek.com/anthropic", Model: "deepseek-chat"} + + toolset := &Toolset{ + Workspace: ".", + Timeout: 30 * time.Second, + MaxOutput: 4000, + apiCfg: apiCfg, + dockerSocket: defaultDockerSocket, + coop: NewCoopManager(), + sessionRoot: t.TempDir(), + } + + // 1. 缺少 task + r := toolset.runCoop(context.Background(), "test-session", map[string]any{}) + if r.Success || r.Output == "" { + t.Fatalf("empty task should fail, got %#v", r) + } + + // 2. 正常参数:应进入 Docker 检查阶段(本机无 Docker 时返回连接/镜像错误而非 panic) + r = toolset.runCoop(context.Background(), "test-session", map[string]any{ + "task": "测试任务", + "round_max": 2, + "timeout": 30, + }) + t.Logf("runCoop output: %s", r.Output) + if r.Success { + // 若真的跑成功了(有 Docker 且镜像存在),也无妨 + t.Logf("coop unexpectedly succeeded") + } +} diff --git a/internal/agent/docker_config.go b/internal/agent/docker_config.go new file mode 100644 index 0000000..026d684 --- /dev/null +++ b/internal/agent/docker_config.go @@ -0,0 +1,99 @@ +package agent + +import ( + "encoding/json" + "os" + "path/filepath" + "runtime" + "sync" +) + +// DockerConfig 是用户可配置的 Docker 连接信息。 +type DockerConfig struct { + // Socket 是用户配置的 Docker 连接地址,支持 + // unix:///var/run/docker.sock、npipe:////./pipe/docker_engine、tcp://host:2375 等。 + // 为空时使用默认的本地 Docker。 + Socket string `json:"socket,omitempty"` +} + +// DockerConfigStore 持久化用户的 Docker socket 配置,并给出默认值。 +type DockerConfigStore struct { + mu sync.RWMutex + path string + config DockerConfig + defaultSocket string +} + +func NewDockerConfigStore(dataDir string) (*DockerConfigStore, error) { + defaultSocket := getenv("AGENT_DOCKER_SOCKET", "") + if defaultSocket == "" { + defaultSocket = defaultDockerSocket() + } + store := &DockerConfigStore{ + path: filepath.Join(dataDir, "docker-config.json"), + defaultSocket: defaultSocket, + } + store.load() + return store, nil +} + +func defaultDockerSocket() string { + if runtime.GOOS == "windows" { + return "npipe:////./pipe/docker_engine" + } + return "unix:///var/run/docker.sock" +} + +func (s *DockerConfigStore) load() { + data, err := os.ReadFile(s.path) + if err != nil { + return + } + if len(data) == 0 { + return + } + var config DockerConfig + if err := json.Unmarshal(data, &config); err != nil { + return + } + s.config = config +} + +// Socket 返回当前生效的 docker 连接地址:用户配置优先,否则使用默认本地 Docker。 +func (s *DockerConfigStore) Socket() string { + s.mu.RLock() + defer s.mu.RUnlock() + if s.config.Socket != "" { + return s.config.Socket + } + return s.defaultSocket +} + +// DefaultSocket 返回未配置时使用的默认地址。 +func (s *DockerConfigStore) DefaultSocket() string { + return s.defaultSocket +} + +// IsConfigured 报告用户是否手动配置过 socket。 +func (s *DockerConfigStore) IsConfigured() bool { + s.mu.RLock() + defer s.mu.RUnlock() + return s.config.Socket != "" +} + +// SetSocket 保存用户的 socket 配置;传入空字符串表示清除配置、恢复默认。 +func (s *DockerConfigStore) SetSocket(socket string) error { + s.mu.Lock() + defer s.mu.Unlock() + + s.config.Socket = socket + data, err := json.MarshalIndent(s.config, "", " ") + if err != nil { + return err + } + tmp := s.path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return err + } + return os.Rename(tmp, s.path) +} diff --git a/internal/agent/llm.go b/internal/agent/llm.go new file mode 100644 index 0000000..e5ed33b --- /dev/null +++ b/internal/agent/llm.go @@ -0,0 +1,470 @@ +package agent + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net/http" + "strings" + "time" +) + +type LLMClient struct { + cfg Config + apiConfig *APIConfigStore + client *http.Client +} + +type StreamChunk struct { + Content string + ToolCalls []ToolCallDelta + FinishReason string + Usage *Usage + Error error +} + +type ToolCallDelta struct { + Index int + ID string + Name string + ArgumentsDelta string +} + +func NewLLMClient(cfg Config, apiCfg *APIConfigStore) *LLMClient { + return &LLMClient{ + cfg: cfg, + apiConfig: apiCfg, + client: &http.Client{ + Timeout: cfg.RequestTimeout + 30*time.Second, + }, + } +} + +// isAnthropic 报告当前是否使用 Anthropic Messages API。 +func (c *LLMClient) isAnthropic() bool { + return c.apiConfig.Provider() == ProviderAnthropic +} + +func (c *LLMClient) ChatStream(ctx context.Context, messages []Message, tools []ToolDefinition) (<-chan StreamChunk, error) { + if c.isAnthropic() { + body := map[string]any{ + "model": c.apiConfig.Model(), + "max_tokens": c.cfg.StreamOutputTokens, + "stream": true, + "temperature": c.cfg.Temperature, + "messages": toAnthropicMessages(messages), + } + if system := extractSystemPrompt(messages); system != "" { + body["system"] = system + } + if len(tools) > 0 { + body["tools"] = toAnthropicTools(tools) + } + return c.stream(ctx, body) + } + + body := map[string]any{ + "model": c.apiConfig.Model(), + "messages": toAPIMessages(messages), + "stream": true, + "temperature": c.cfg.Temperature, + "max_tokens": c.cfg.StreamOutputTokens, + } + if len(tools) > 0 { + body["tools"] = tools + body["tool_choice"] = "auto" + } + return c.stream(ctx, body) +} + +func (c *LLMClient) Compress(ctx context.Context, history []Message) (string, error) { + systemPrompt := `你是一个高信息密度的对话压缩器。把下面的历史对话压缩成一份中文摘要,供后续继续执行同一任务。 +必须保留(按重要程度排序): +1. 用户的核心目标与最新要求、尚未完成的任务与下一步计划; +2. 已确认的决策与结论、得分/进度(如 TSec 累计分数与已通关/进行中的题目编号); +3. 关键文件路径、执行过的命令与重要输出要点、代码要点; +4. 错误信息与解决方案、需要继续跟进的问题; +5. 平台提交规则(如 BENCHMARK_TOKEN、unique_code、提交接口与提交方式)。 +规则: +- 按时间顺序组织,最新信息优先,可适当合并同类项; +- 若历史包含工具调用(run_bash/read_file/run_coop 等),只保留"做了什么、结果如何"的要点,不要逐字复制命令或输出; +- 不要添加历史中不存在的信息,不要臆测; +- 直接输出摘要正文,不要输出任何解释、标题或多余格式。` + messages := append([]Message{{Role: "system", Content: &systemPrompt}}, history...) + + if c.isAnthropic() { + return c.compressAnthropic(ctx, messages) + } + + body := map[string]any{ + "model": c.apiConfig.Model(), + "messages": toAPIMessages(messages), + "stream": false, + "temperature": 0.2, + "max_tokens": c.cfg.CompactionTokens, + } + + ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout) + defer cancel() + + payload, err := json.Marshal(body) + if err != nil { + return "", err + } + respBody, err := c.doSyncRequestWithRetry(ctx, false, c.apiConfig.BaseURL(), payload) + if err != nil { + return "", err + } + + var result struct { + Choices []struct { + Message struct { + Content *string `json:"content"` + ReasoningContent *string `json:"reasoning_content"` + } `json:"message"` + } `json:"choices"` + } + if err := json.Unmarshal(respBody, &result); err != nil { + return "", err + } + if len(result.Choices) == 0 { + return "", errors.New("模型没有返回压缩摘要") + } + msg := result.Choices[0].Message + if msg.Content != nil && strings.TrimSpace(*msg.Content) != "" { + return strings.TrimSpace(*msg.Content), nil + } + // 推理模型(deepseek-reasoner 系)正文在 reasoning_content,content 可能为空 + if msg.ReasoningContent != nil && strings.TrimSpace(*msg.ReasoningContent) != "" { + return strings.TrimSpace(*msg.ReasoningContent), nil + } + return "", errors.New("模型没有返回压缩摘要") +} + +// stream 发起流式 LLM 请求,对连接失败、限流(429/5xx/551)、首 chunk 前流断开 +// 做指数退避重试(3 次:1s→2s→4s)。保留 (<-chan, error) 签名,上层无需改动。 +func (c *LLMClient) stream(ctx context.Context, body map[string]any) (<-chan StreamChunk, error) { + anthropic := c.isAnthropic() + url := c.apiConfig.BaseURL() + if anthropic { + url = normalizeAnthropicURL(url) + } + payload, err := json.Marshal(body) + if err != nil { + return nil, err + } + + const maxRetries = 3 + var lastErr error + + for attempt := 0; attempt <= maxRetries; attempt++ { + if attempt > 0 { + backoff := time.Duration(1<<(attempt-1)) * time.Second + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(backoff): + } + } + + ch, reqErr := c.doStreamRequest(ctx, anthropic, url, payload) + if reqErr != nil { + lastErr = reqErr + if attempt < maxRetries && isRetriableErr(reqErr) { + log.Printf("[llm] 流式请求失败(第 %d 次),将重试: %v", attempt+1, reqErr) + continue + } + return nil, reqErr + } + + // 请求成功建立,但流可能在首个 chunk 前就断开(EOF)。 + // peek 第一个 chunk:如果是可重试错误,排空旧 channel 后重试。 + first, ok := <-ch + if !ok { + lastErr = errors.New("stream: 空响应") + if attempt < maxRetries { + log.Printf("[llm] 流式响应为空(第 %d 次),重试", attempt+1) + continue + } + return nil, lastErr + } + if first.Error != nil && attempt < maxRetries && isRetriableErr(first.Error) { + for range ch { + } + lastErr = first.Error + log.Printf("[llm] 流式响应首 chunk 前断开(第 %d 次),重试: %v", attempt+1, first.Error) + continue + } + + // 正常:转发 first + 后续 chunk + out := make(chan StreamChunk, 64) + go func() { + defer close(out) + out <- first + for chunk := range ch { + out <- chunk + } + }() + return out, nil + } + + return nil, lastErr +} + +// doStreamRequest 发起单次流式 LLM 请求(不含重试)。 +func (c *LLMClient) doStreamRequest(ctx context.Context, anthropic bool, url string, payload []byte) (<-chan StreamChunk, error) { + ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout) + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + cancel() + return nil, err + } + c.setRequestHeaders(req, anthropic) + + resp, err := c.client.Do(req) + if err != nil { + cancel() + return nil, err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + apiErr := c.readAPIError(resp) + resp.Body.Close() + cancel() + return nil, apiErr + } + + ch := make(chan StreamChunk, 64) + go func() { + defer cancel() + defer close(ch) + defer resp.Body.Close() + if anthropic { + readAnthropicSSE(ctx, resp, ch) + } else { + c.readSSE(ctx, resp, ch) + } + }() + return ch, nil +} + +// doSyncRequestWithRetry 发起同步(非流式)LLM 请求,对连接失败、限流(429/5xx/551) +// 做指数退避重试(3 次:1s→2s→4s)。返回响应体字节,由调用方解析。 +func (c *LLMClient) doSyncRequestWithRetry(ctx context.Context, anthropic bool, url string, payload []byte) ([]byte, error) { + const maxRetries = 3 + var lastErr error + + for attempt := 0; attempt <= maxRetries; attempt++ { + if attempt > 0 { + backoff := time.Duration(1<<(attempt-1)) * time.Second + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(backoff): + } + } + + body, err := c.doSyncRequest(ctx, anthropic, url, payload) + if err == nil { + return body, nil + } + lastErr = err + if attempt < maxRetries && isRetriableErr(err) { + log.Printf("[llm] 同步请求失败(第 %d 次),将重试: %v", attempt+1, err) + continue + } + return nil, err + } + return nil, lastErr +} + +// doSyncRequest 发起单次同步(非流式)LLM 请求(不含重试)。 +func (c *LLMClient) doSyncRequest(ctx context.Context, anthropic bool, url string, payload []byte) ([]byte, error) { + ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(payload)) + if err != nil { + return nil, err + } + c.setRequestHeaders(req, anthropic) + + resp, err := c.client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, c.readAPIError(resp) + } + return io.ReadAll(resp.Body) +} + +// isRetriableErr 判断错误是否值得重试(网络错误、限流、服务端错误)。 +func isRetriableErr(err error) bool { + if err == nil { + return false + } + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) { + return true + } + msg := err.Error() + if strings.Contains(msg, "EOF") || + strings.Contains(msg, "connection reset") || + strings.Contains(msg, "broken pipe") || + strings.Contains(msg, "timeout") || + strings.Contains(msg, "deadline exceeded") || + strings.Contains(msg, "connection refused") || + strings.Contains(msg, "no such host") { + return true + } + // HTTP 状态码错误:5xx 服务端错误、429 限流、551 网关熔断 + if strings.HasPrefix(msg, "HTTP 5") || + strings.HasPrefix(msg, "HTTP 429") || + strings.HasPrefix(msg, "HTTP 551") { + return true + } + return false +} + +func (c *LLMClient) readSSE(ctx context.Context, resp *http.Response, ch chan<- StreamChunk) { + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 64*1024), 1024*1024) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "[DONE]" { + break + } + + var chunk struct { + Choices []struct { + Delta struct { + Content *string `json:"content"` + ToolCalls []struct { + Index *int `json:"index"` + ID string `json:"id"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` + } `json:"tool_calls"` + } `json:"delta"` + FinishReason *string `json:"finish_reason"` + } `json:"choices"` + Usage *Usage `json:"usage"` + } + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + continue + } + + event := StreamChunk{} + for _, choice := range chunk.Choices { + if choice.Delta.Content != nil { + event.Content += *choice.Delta.Content + } + for _, tc := range choice.Delta.ToolCalls { + index := 0 + if tc.Index != nil { + index = *tc.Index + } + event.ToolCalls = append(event.ToolCalls, ToolCallDelta{ + Index: index, + ID: tc.ID, + Name: tc.Function.Name, + ArgumentsDelta: tc.Function.Arguments, + }) + } + if choice.FinishReason != nil { + event.FinishReason = *choice.FinishReason + } + } + if chunk.Usage != nil { + event.Usage = chunk.Usage + } + if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" || event.Usage != nil { + ch <- event + } + } + if err := scanner.Err(); err != nil && ctx.Err() == nil { + ch <- StreamChunk{Error: err} + } +} + +func (c *LLMClient) setRequestHeaders(req *http.Request, anthropic bool) { + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + if anthropic { + req.Header.Set("x-api-key", c.apiConfig.APIKey()) + req.Header.Set("anthropic-version", "2023-06-01") + } else { + req.Header.Set("Authorization", "Bearer "+c.apiConfig.APIKey()) + } +} + +func (c *LLMClient) readAPIError(resp *http.Response) error { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 8*1024)) + var apiErr struct { + Error struct { + Message string `json:"message"` + } `json:"error"` + } + _ = json.Unmarshal(body, &apiErr) + message := strings.TrimSpace(apiErr.Error.Message) + if message == "" { + message = strings.TrimSpace(string(body)) + } + if message == "" { + message = fmt.Sprintf("HTTP %d", resp.StatusCode) + } + return errors.New(message) +} + +func toAPIMessages(messages []Message) []map[string]any { + result := make([]map[string]any, 0, len(messages)) + for _, message := range messages { + item := map[string]any{ + "role": message.Role, + "content": contentValue(message.Content), + } + if len(message.ToolCalls) > 0 { + item["tool_calls"] = toAPIToolCalls(message.ToolCalls) + } + if message.ToolCallID != "" { + item["tool_call_id"] = message.ToolCallID + } + result = append(result, item) + } + return result +} + +func toAPIToolCalls(calls []ToolCall) []map[string]any { + result := make([]map[string]any, 0, len(calls)) + for _, call := range calls { + result = append(result, map[string]any{ + "id": call.ID, + "type": "function", + "function": map[string]any{ + "name": call.Function.Name, + "arguments": call.Function.Arguments, + }, + }) + } + return result +} + +func contentValue(content *string) any { + if content == nil { + return nil + } + return *content +} diff --git a/internal/agent/llm_anthropic.go b/internal/agent/llm_anthropic.go new file mode 100644 index 0000000..1983be5 --- /dev/null +++ b/internal/agent/llm_anthropic.go @@ -0,0 +1,298 @@ +package agent + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "net/http" + "strings" +) + +// normalizeAnthropicURL 把用户填写的 Base URL 规范化为 Anthropic Messages 接口地址。 +// 兼容多种填写方式:https://api.anthropic.com、.../v1、.../v1/messages、.../v1/chat/completions。 +func normalizeAnthropicURL(baseURL string) string { + trimmed := strings.TrimRight(strings.TrimSpace(baseURL), "/") + if trimmed == "" { + return "https://api.anthropic.com/v1/messages" + } + switch { + case strings.HasSuffix(trimmed, "/v1/chat/completions"): + return strings.TrimSuffix(trimmed, "/chat/completions") + "/messages" + case strings.HasSuffix(trimmed, "/v1/messages"): + return trimmed + case strings.HasSuffix(trimmed, "/v1"): + return trimmed + "/messages" + default: + return trimmed + "/v1/messages" + } +} + +// extractSystemPrompt 汇总消息中的 system 角色内容,Anthropic 要求 system 放在顶层字段。 +func extractSystemPrompt(messages []Message) string { + var parts []string + for _, message := range messages { + if message.Role == "system" && message.Content != nil && strings.TrimSpace(*message.Content) != "" { + parts = append(parts, *message.Content) + } + } + return strings.Join(parts, "\n\n") +} + +// toAnthropicTools 把内部工具定义转换为 Anthropic 的 tools 数组(input_schema 替代 parameters)。 +func toAnthropicTools(tools []ToolDefinition) []map[string]any { + result := make([]map[string]any, 0, len(tools)) + for _, tool := range tools { + schema := tool.Function.Parameters + if schema == nil { + schema = map[string]any{"type": "object"} + } + result = append(result, map[string]any{ + "name": tool.Function.Name, + "description": tool.Function.Description, + "input_schema": schema, + }) + } + return result +} + +// toAnthropicMessages 把内部 Message 列表转换为 Anthropic messages 数组。 +// system 消息被过滤(走顶层 system 字段);assistant 的 tool_use 与 user 的 +// tool_result 都以 content block 形式表达。 +// 注意:Anthropic 要求上一条 assistant 消息中所有 tool_use 的 tool_result +// 必须放在紧邻的同一条 user 消息里,因此连续的 tool 结果消息需要合并。 +func toAnthropicMessages(messages []Message) []map[string]any { + result := make([]map[string]any, 0, len(messages)) + for i := 0; i < len(messages); i++ { + message := messages[i] + switch message.Role { + case "system": + continue + case "assistant": + blocks := make([]map[string]any, 0, 1+len(message.ToolCalls)) + if message.Content != nil && strings.TrimSpace(*message.Content) != "" { + blocks = append(blocks, map[string]any{"type": "text", "text": *message.Content}) + } + for _, call := range message.ToolCalls { + blocks = append(blocks, map[string]any{ + "type": "tool_use", + "id": call.ID, + "name": call.Function.Name, + "input": parseJSONValue(call.Function.Arguments), + }) + } + if len(blocks) == 0 { + blocks = append(blocks, map[string]any{"type": "text", "text": ""}) + } + result = append(result, map[string]any{"role": "assistant", "content": blocks}) + case "tool": + // 合并连续的 tool 消息:同一条 user 消息包含所有 tool_result 块 + blocks := []map[string]any{{ + "type": "tool_result", + "tool_use_id": message.ToolCallID, + "content": contentString(message.Content), + }} + for i+1 < len(messages) && messages[i+1].Role == "tool" { + i++ + next := messages[i] + blocks = append(blocks, map[string]any{ + "type": "tool_result", + "tool_use_id": next.ToolCallID, + "content": contentString(next.Content), + }) + } + result = append(result, map[string]any{"role": "user", "content": blocks}) + default: // user + blocks := make([]map[string]any, 0, 1) + if message.Content != nil && strings.TrimSpace(*message.Content) != "" { + blocks = append(blocks, map[string]any{"type": "text", "text": *message.Content}) + } + if len(blocks) == 0 { + blocks = append(blocks, map[string]any{"type": "text", "text": ""}) + } + result = append(result, map[string]any{"role": "user", "content": blocks}) + } + } + return result +} + +// parseJSONValue 把工具参数 JSON 字符串解析为任意值;解析失败时回退为空对象。 +func parseJSONValue(raw string) any { + var value any + if err := json.Unmarshal([]byte(raw), &value); err != nil || value == nil { + return map[string]any{} + } + return value +} + +func contentString(content *string) string { + if content == nil { + return "" + } + return *content +} + +// readAnthropicSSE 解析 Anthropic Messages API 的流式响应。 +// 事件格式为 `event: ` 与 `data: ` 两行一组。 +func readAnthropicSSE(ctx context.Context, resp *http.Response, ch chan<- StreamChunk) { + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 64*1024), 1024*1024) + + // content block index -> 正在累积的 tool_use 状态 + type toolState struct { + callIndex int + id string + name string + } + tools := make(map[int]*toolState) + nextCallIndex := 0 + + var eventType string + for scanner.Scan() { + line := scanner.Text() + switch { + case strings.HasPrefix(line, "event:"): + eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + continue + case strings.HasPrefix(line, "data:"): + default: + continue + } + + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + switch eventType { + case "content_block_start": + var ev struct { + Index int `json:"index"` + Block struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + } `json:"content_block"` + } + if err := json.Unmarshal([]byte(data), &ev); err != nil { + continue + } + if ev.Block.Type == "tool_use" { + state := &toolState{callIndex: nextCallIndex, id: ev.Block.ID, name: ev.Block.Name} + nextCallIndex++ + tools[ev.Index] = state + ch <- StreamChunk{ToolCalls: []ToolCallDelta{ + {Index: state.callIndex, ID: state.id, Name: state.name}, + }} + } + case "content_block_delta": + var ev struct { + Index int `json:"index"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + PartialJSON string `json:"partial_json"` + } `json:"delta"` + } + if err := json.Unmarshal([]byte(data), &ev); err != nil { + continue + } + switch ev.Delta.Type { + case "text_delta": + ch <- StreamChunk{Content: ev.Delta.Text} + case "input_json_delta": + if state, ok := tools[ev.Index]; ok { + ch <- StreamChunk{ToolCalls: []ToolCallDelta{ + {Index: state.callIndex, ArgumentsDelta: ev.Delta.PartialJSON}, + }} + } + } + case "message_delta": + var ev struct { + Delta struct { + StopReason string `json:"stop_reason"` + } `json:"delta"` + Usage *struct { + OutputTokens int `json:"output_tokens"` + } `json:"usage"` + } + if err := json.Unmarshal([]byte(data), &ev); err != nil { + continue + } + if ev.Delta.StopReason != "" { + ch <- StreamChunk{FinishReason: ev.Delta.StopReason} + } + if ev.Usage != nil { + ch <- StreamChunk{Usage: &Usage{CompletionTokens: ev.Usage.OutputTokens}} + } + case "error": + var ev struct { + Error struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal([]byte(data), &ev); err != nil { + continue + } + msg := strings.TrimSpace(ev.Error.Message) + if msg == "" { + msg = "Anthropic API 错误" + } + ch <- StreamChunk{Error: errors.New(msg)} + } + } + if err := scanner.Err(); err != nil && ctx.Err() == nil { + ch <- StreamChunk{Error: err} + } +} + +// compressAnthropic 使用 Anthropic 非流式接口执行对话压缩。 +func (c *LLMClient) compressAnthropic(ctx context.Context, messages []Message) (string, error) { + body := map[string]any{ + "model": c.apiConfig.Model(), + "max_tokens": c.cfg.CompactionTokens, + "stream": false, + "temperature": 0.2, + // 关闭思考:推理模型(如 deepseek-v4-flash)默认先输出 thinking 块, + // 会把 max_tokens 预算耗尽而拿不到 text 块,导致压缩被判为失败。 + "thinking": map[string]any{"type": "disabled"}, + "messages": toAnthropicMessages(messages), + } + if system := extractSystemPrompt(messages); system != "" { + body["system"] = system + } + + ctx, cancel := context.WithTimeout(ctx, c.cfg.RequestTimeout) + defer cancel() + + payload, err := json.Marshal(body) + if err != nil { + return "", err + } + respBody, err := c.doSyncRequestWithRetry(ctx, true, normalizeAnthropicURL(c.apiConfig.BaseURL()), payload) + if err != nil { + return "", err + } + + var result struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + Thinking string `json:"thinking"` + } `json:"content"` + } + if err := json.Unmarshal(respBody, &result); err != nil { + return "", err + } + fallback := "" + for _, block := range result.Content { + if block.Type == "text" && strings.TrimSpace(block.Text) != "" { + return strings.TrimSpace(block.Text), nil + } + // 记录 thinking 作为兜底(仅当端点不支持 thinking:disabled 时才会出现) + if block.Type == "thinking" && fallback == "" { + fallback = strings.TrimSpace(block.Thinking) + } + } + if fallback != "" { + return fallback, nil + } + return "", errors.New("模型没有返回压缩摘要") +} diff --git a/internal/agent/session.go b/internal/agent/session.go new file mode 100644 index 0000000..0abc1af --- /dev/null +++ b/internal/agent/session.go @@ -0,0 +1,338 @@ +package agent + +import ( + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + "unicode/utf8" +) + +type SessionStore struct { + mu sync.RWMutex + path string + sessions map[string]*Session + order []string +} + +func NewSessionStore(dataDir string) (*SessionStore, error) { + if err := os.MkdirAll(dataDir, 0o755); err != nil { + return nil, err + } + path := filepath.Join(dataDir, "sessions.json") + store := &SessionStore{ + path: path, + sessions: make(map[string]*Session), + order: make([]string, 0), + } + _ = store.load() + return store, nil +} + +func (s *SessionStore) load() error { + data, err := os.ReadFile(s.path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err + } + if len(strings.TrimSpace(string(data))) == 0 { + return nil + } + var sessions map[string]*Session + if err := json.Unmarshal(data, &sessions); err != nil { + return err + } + for id, session := range sessions { + if session == nil || session.ID == "" { + continue + } + s.sessions[id] = session + s.order = append(s.order, id) + } + s.sortOrder() + return nil +} + +func (s *SessionStore) saveLocked() error { + data, err := json.MarshalIndent(s.sessions, "", " ") + if err != nil { + return err + } + tmp := s.path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return err + } + return os.Rename(tmp, s.path) +} + +func (s *SessionStore) List() []SessionSummary { + s.mu.RLock() + defer s.mu.RUnlock() + return s.listLocked("", 0, 0) +} + +// Search 按关键词过滤会话(标题或任意消息内容,忽略大小写)。 +func (s *SessionStore) Search(query string) []SessionSummary { + s.mu.RLock() + defer s.mu.RUnlock() + return s.listLocked(query, 0, 0) +} + +// listLocked 返回会话摘要列表;query 非空时按关键词过滤,limit>0 时截断。 +// 必须持读锁调用。 +func (s *SessionStore) listLocked(query string, limit, offset int) []SessionSummary { + query = strings.ToLower(strings.TrimSpace(query)) + out := make([]SessionSummary, 0, len(s.order)) + skip := 0 + for _, id := range s.order { + session, ok := s.sessions[id] + if !ok { + continue + } + if query != "" && !sessionMatches(session, query) { + continue + } + if offset > 0 && skip < offset { + skip++ + continue + } + if limit > 0 && len(out) >= limit { + break + } + out = append(out, summarize(session)) + } + return out +} + +func sessionMatches(session *Session, query string) bool { + if strings.Contains(strings.ToLower(session.Title), query) { + return true + } + for _, message := range session.Messages { + if message.Content != nil && strings.Contains(strings.ToLower(*message.Content), query) { + return true + } + for _, call := range message.ToolCalls { + if strings.Contains(strings.ToLower(call.Function.Name), query) || + strings.Contains(strings.ToLower(call.Function.Arguments), query) { + return true + } + } + } + return false +} + +// Messages 分页返回会话消息(按时间正序),offset 从最新一条往前数, +// 即 offset=0 返回最新 limit 条;返回 has_more 表示还有更早的消息。 +func (s *SessionStore) Messages(id string, limit, offset int) ([]Message, bool, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + session, ok := s.sessions[id] + if !ok { + return nil, false, false + } + total := len(session.Messages) + if limit <= 0 { + limit = 50 + } + if offset < 0 { + offset = 0 + } + start := total - offset - limit + if start < 0 { + start = 0 + } + messages := make([]Message, 0, limit) + for i := start; i < total-offset; i++ { + messages = append(messages, cloneMessage(session.Messages[i])) + } + hasMore := start > 0 + return messages, true, hasMore +} + +// Clear 清空会话消息(保留会话本身)。 +func (s *SessionStore) Clear(id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + session, ok := s.sessions[id] + if !ok { + return false + } + session.Messages = make([]Message, 0) + session.UpdatedAt = time.Now() + _ = s.saveLocked() + return true +} + +// TotalMessages 返回会话消息总数(用于分页计算)。 +func (s *SessionStore) TotalMessages(id string) (int, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + session, ok := s.sessions[id] + if !ok { + return 0, false + } + return len(session.Messages), true +} + +func cloneMessage(message Message) Message { + copy := message + if message.Content != nil { + content := *message.Content + copy.Content = &content + } + copy.ToolCalls = append([]ToolCall(nil), message.ToolCalls...) + return copy +} + +func (s *SessionStore) Get(id string) (*Session, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + session, ok := s.sessions[id] + return session, ok +} + +func (s *SessionStore) Create() *Session { + s.mu.Lock() + defer s.mu.Unlock() + + session := &Session{ + ID: newID(), + Title: "新对话", + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + Messages: make([]Message, 0), + } + s.sessions[session.ID] = session + s.order = append([]string{session.ID}, s.order...) + _ = s.saveLocked() + return session +} + +func (s *SessionStore) Delete(id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.sessions[id]; !ok { + return false + } + delete(s.sessions, id) + for i, item := range s.order { + if item == id { + s.order = append(s.order[:i], s.order[i+1:]...) + break + } + } + _ = s.saveLocked() + return true +} + +func (s *SessionStore) Save(session *Session) { + session.UpdatedAt = time.Now() + s.mu.Lock() + // 存储克隆而非原对象:会话对象可能仍在 Agent 循环中被修改, + // 直接存指针会导致 saveLocked 序列化时与其他会话的写入产生数据竞争。 + s.sessions[session.ID] = cloneSession(session) + s.order = append([]string{session.ID}, removeString(s.order, session.ID)...) + err := s.saveLocked() + s.mu.Unlock() + if err != nil { + // A failed save should not break an in-memory conversation. + return + } +} + +func (s *SessionStore) sortOrder() { + sort.SliceStable(s.order, func(i, j int) bool { + a, aOK := s.sessions[s.order[i]] + b, bOK := s.sessions[s.order[j]] + if !aOK || !bOK { + return aOK && !bOK + } + return a.UpdatedAt.After(b.UpdatedAt) + }) +} + +func summarize(session *Session) SessionSummary { + summary := SessionSummary{ + ID: session.ID, + Title: session.Title, + CreatedAt: session.CreatedAt, + UpdatedAt: session.UpdatedAt, + MessageCount: len(session.Messages), + } + for i := len(session.Messages) - 1; i >= 0; i-- { + message := session.Messages[i] + if message.Role != "user" && message.Role != "assistant" { + continue + } + if message.Content != nil { + summary.Preview = truncateRunes(*message.Content, 120) + break + } + } + return summary +} + +func newID() string { + buf := make([]byte, 12) + _, _ = rand.Read(buf) + return hex.EncodeToString(buf) +} + +// cloneSession 深拷贝会话对象,保证持久化对象在存储后不再被外部修改, +// 从而消除 Agent 循环写入与 saveLocked 序列化之间的数据竞争。 +func cloneSession(session *Session) *Session { + if session == nil { + return nil + } + clone := *session + clone.Messages = make([]Message, len(session.Messages)) + for i, message := range session.Messages { + copy := message + if message.Content != nil { + content := *message.Content + copy.Content = &content + } + copy.ToolCalls = append([]ToolCall(nil), message.ToolCalls...) + clone.Messages[i] = copy + } + return &clone +} + +func removeString(items []string, target string) []string { + result := make([]string, 0, len(items)) + for _, item := range items { + if item != target { + result = append(result, item) + } + } + return result +} + +func truncateRunes(value string, max int) string { + runes := []rune(value) + if len(runes) <= max { + return value + } + return string(runes[:max]) + "..." +} + +func firstRunes(value string, max int) string { + runes := []rune(strings.TrimSpace(value)) + if len(runes) <= max { + return string(runes) + } + return string(runes[:max]) +} + +func countRunes(value string) int { + return utf8.RuneCountInString(value) +} diff --git a/internal/agent/tools.go b/internal/agent/tools.go new file mode 100644 index 0000000..8ae12fd --- /dev/null +++ b/internal/agent/tools.go @@ -0,0 +1,559 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + "time" +) + +type Toolset struct { + Workspace string + Timeout time.Duration + MaxOutput int + shell string + shellArgs []string + python string + dockerSocket func() string + apiCfg *APIConfigStore + coop *CoopManager + // sessionRoot 是各会话独立工作目录的根(data/workspaces), + // 每个会话在其下拥有 / 目录,互不干扰。 + sessionRoot string +} + +func NewToolset(cfg Config, dockerCfg *DockerConfigStore, apiCfg *APIConfigStore) *Toolset { + toolset := &Toolset{ + Workspace: cfg.Workspace, + Timeout: cfg.ToolTimeout, + MaxOutput: cfg.MaxToolResultChars, + apiCfg: apiCfg, + sessionRoot: filepath.Join(cfg.DataDir, "workspaces"), + } + + if dockerCfg != nil { + toolset.dockerSocket = dockerCfg.Socket + } else { + toolset.dockerSocket = defaultDockerSocket + } + + if shell, err := exec.LookPath("bash"); err == nil { + toolset.shell = shell + toolset.shellArgs = []string{"-lc"} + } else if shell, err := exec.LookPath("pwsh"); err == nil { + toolset.shell = shell + toolset.shellArgs = []string{"-NoProfile", "-NonInteractive", "-Command"} + } else if shell, err := exec.LookPath("powershell"); err == nil { + toolset.shell = shell + toolset.shellArgs = []string{"-NoProfile", "-NonInteractive", "-Command"} + } else { + toolset.shell = "cmd.exe" + toolset.shellArgs = []string{"/C"} + } + + if python, err := exec.LookPath("python"); err == nil { + toolset.python = python + } else if python, err := exec.LookPath("python3"); err == nil { + toolset.python = python + } else { + toolset.python = "python" + } + return toolset +} + +func (t *Toolset) Definitions() []ToolDefinition { + return []ToolDefinition{ + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "run_bash", + Description: "在用户工作区目录中执行一条 Bash 命令。Windows 上没有 Bash 时自动回退到 PowerShell 或 cmd。适合查看文件、运行构建、安装依赖、搜索代码、启动程序等。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "command": map[string]any{ + "type": "string", + "description": "要执行的完整命令,例如 `ls -la` 或 `go test ./...`。", + }, + "timeout_seconds": map[string]any{ + "type": "integer", + "description": "命令超时秒数,默认 120 秒。", + }, + }, + "required": []string{"command"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "read_file", + Description: "读取工作区内的文本文件并返回内容。超长文件会被截断。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "path": map[string]any{ + "type": "string", + "description": "文件路径,可以是相对路径或绝对路径。必须位于工作区内。", + }, + }, + "required": []string{"path"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "write_file", + Description: "写入或覆盖工作区内的文件。父目录不存在时自动创建。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "path": map[string]any{ + "type": "string", + "description": "目标文件路径,可以是相对路径或绝对路径。必须位于工作区内。", + }, + "content": map[string]any{ + "type": "string", + "description": "要写入的完整文件内容。", + }, + }, + "required": []string{"path", "content"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "run_python", + Description: "在工作区目录中执行一段 Python 代码。适合数据处理、批量修改、生成脚本和自动化任务。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "code": map[string]any{ + "type": "string", + "description": "要执行的完整 Python 代码。", + }, + }, + "required": []string{"code"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "list_directory", + Description: "列出工作区内目录的内容,包括文件大小和修改时间。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "path": map[string]any{ + "type": "string", + "description": "目录路径,默认为工作区根目录。", + }, + }, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_ps", + Description: "列出 Docker 容器。默认连接本地 Docker,可通过 Web 设置中的 Docker Socket 切换到远程 Docker 服务器。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "all": map[string]any{ + "type": "boolean", + "description": "是否列出所有容器(包括已停止的),默认 false 只显示运行中。", + }, + }, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_images", + Description: "列出 Docker 镜像列表。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_logs", + Description: "查看 Docker 容器的最近日志。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "container": map[string]any{ + "type": "string", + "description": "容器名称或 ID。", + }, + "tail": map[string]any{ + "type": "integer", + "description": "返回最近的日志行数,默认 100。", + }, + }, + "required": []string{"container"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_inspect", + Description: "查看 Docker 容器的详细配置和状态信息。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "container": map[string]any{ + "type": "string", + "description": "容器名称或 ID。", + }, + }, + "required": []string{"container"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_exec", + Description: "在 Docker 容器内执行一条命令(使用 sh -c),返回执行输出。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "container": map[string]any{ + "type": "string", + "description": "容器名称或 ID。", + }, + "command": map[string]any{ + "type": "string", + "description": "要在容器内执行的命令,例如 `ls -la` 或 `cat /etc/os-release`。", + }, + }, + "required": []string{"container", "command"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_start", + Description: "启动一个 Docker 容器。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "container": map[string]any{ + "type": "string", + "description": "容器名称或 ID。", + }, + }, + "required": []string{"container"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "docker_stop", + Description: "停止一个 Docker 容器。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "container": map[string]any{ + "type": "string", + "description": "容器名称或 ID。", + }, + "timeout": map[string]any{ + "type": "integer", + "description": "等待优雅停止的秒数,默认由 Docker 决定。", + }, + }, + "required": []string{"container"}, + }, + }, + }, + { + Type: "function", + Function: ToolDefinitionFunction{ + Name: "run_coop", + Description: "调用协作容器异步完成一个任务(容器内单个 agent 独立执行,自动注入当前 Agent 已配置的模型、Base URL、API Key 与接口协议,无需重复填写)。支持三种引擎:engine=pi(默认,pi-coding-agent,支持 openai+anthropic 协议)、engine=pigo(pigo,支持 openai+anthropic 协议)或 engine=claude(Claude Code,仅支持 anthropic 协议)。需先构建对应镜像(pi-coop / pigo-coop / claude-coop)。该工具是异步的:调用后立即返回任务 ID,容器在后台运行;任务完成(成功 / 失败 / 超时)后系统会自动注入一条【协作任务完成通知】消息(含结构化结果摘要),由你确认结果、必要时提交 flag 并向用户汇报,无需在此等待。", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "task": map[string]any{ + "type": "string", + "description": "要交给协作容器完成的任务描述,例如“请为 XX 编写设计文档并实现原型”。解题类任务需写明目标地址、平台提交规则(BENCHMARK_TOKEN / unique_code / submit API)。", + }, + "engine": map[string]any{ + "type": "string", + "description": "协作引擎:pi(默认,pi-coding-agent)、pigo(pigo)或 claude(Claude Code,仅支持 anthropic 协议)。pi 与 pigo 支持 openai+anthropic 协议;claude 仅支持 anthropic 协议。", + "enum": []string{"pi", "pigo", "claude"}, + }, + "round_max": map[string]any{ + "type": "integer", + "description": "最大轮次,默认 1(单 agent 一次运行完成,未完成则失败并重新下发)。", + }, + "timeout": map[string]any{ + "type": "integer", + "description": "单轮超时秒数,默认 1800(30 分钟)。单 agent 需在此时限内完成全部工作并创建 DONE。", + }, + "blackboard": map[string]any{ + "type": "string", + "description": "可选:主机上的目录路径,挂载到容器 /blackboard 保留黑板产物与会话;Windows 路径会自动转换为 WSL 挂载路径(/mnt/盘符/...)。留空则自动使用本会话工作区下的独立子目录 blackboard/<任务ID>(同一会话多次协作互不污染)。", + }, + }, + "required": []string{"task"}, + }, + }, + }, + } +} + +func (t *Toolset) Execute(ctx context.Context, sessionID, name, arguments string) ToolResult { + var params map[string]any + if err := json.Unmarshal([]byte(arguments), ¶ms); err != nil { + return ToolResult{Success: false, Output: "无法解析工具参数: " + err.Error()} + } + + switch name { + case "run_bash": + command, _ := params["command"].(string) + if strings.TrimSpace(command) == "" { + return ToolResult{Success: false, Output: "command 不能为空"} + } + return t.runShell(ctx, sessionID, command, intParam(params, "timeout_seconds")) + case "read_file": + path, _ := params["path"].(string) + return t.readFile(sessionID, path) + case "write_file": + path, _ := params["path"].(string) + content, _ := params["content"].(string) + return t.writeFile(sessionID, path, content) + case "run_python": + code, _ := params["code"].(string) + if strings.TrimSpace(code) == "" { + return ToolResult{Success: false, Output: "code 不能为空"} + } + return t.runPython(ctx, sessionID, code) + case "list_directory": + path, _ := params["path"].(string) + return t.listDirectory(sessionID, path) + case "docker_ps": + return t.dockerPS(ctx, params) + case "docker_images": + return t.dockerImages(ctx, params) + case "docker_logs": + return t.dockerLogs(ctx, params) + case "docker_inspect": + return t.dockerInspect(ctx, params) + case "docker_exec": + return t.dockerExec(ctx, params) + case "docker_start": + return t.dockerStart(ctx, params) + case "docker_stop": + return t.dockerStop(ctx, params) + case "run_coop": + return t.runCoop(ctx, sessionID, params) + default: + return ToolResult{Success: false, Output: "未知工具: " + name} + } +} + +func (t *Toolset) runShell(ctx context.Context, sessionID, command string, timeoutSeconds int) ToolResult { + // 不信任 LLM 传入的超长 timeout:超过全局工具超时(默认 120s)一律截断为全局值。 + // 否则 sleep / 轮询类命令会长时间阻塞主 Agent(持有会话锁),期间前端 + // 长时间收不到事件会被判"任务中断",且 coop 完成通知也拿不到锁无法及时处理。 + if timeoutSeconds <= 0 || timeoutSeconds > int(t.Timeout.Seconds()) { + timeoutSeconds = int(t.Timeout.Seconds()) + } + ctx, cancel := context.WithTimeout(ctx, time.Duration(timeoutSeconds)*time.Second) + defer cancel() + + var stdout, stderr bytes.Buffer + cmd := exec.CommandContext(ctx, t.shell, t.shellArgs...) + cmd.Args = append(cmd.Args, command) + cmd.Dir = t.workspaceFor(sessionID) + cmd.Env = append(os.Environ(), "AGENT_WORKSPACE="+t.workspaceFor(sessionID)) + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + var out strings.Builder + out.WriteString(strings.TrimSpace(stdout.String())) + if stderrText := strings.TrimSpace(stderr.String()); stderrText != "" { + if out.Len() > 0 { + out.WriteString("\n") + } + out.WriteString("[stderr]\n") + out.WriteString(stderrText) + } + if err != nil { + if out.Len() > 0 { + out.WriteString("\n") + } + out.WriteString("[命令失败] " + err.Error()) + } + return ToolResult{Success: err == nil, Output: t.truncate(out.String())} +} + +func (t *Toolset) runPython(ctx context.Context, sessionID, code string) ToolResult { + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + + var stdout, stderr bytes.Buffer + cmd := exec.CommandContext(ctx, t.python, "-c", code) + cmd.Dir = t.workspaceFor(sessionID) + cmd.Env = append(os.Environ(), "AGENT_WORKSPACE="+t.workspaceFor(sessionID)) + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + var out strings.Builder + out.WriteString(strings.TrimSpace(stdout.String())) + if stderrText := strings.TrimSpace(stderr.String()); stderrText != "" { + if out.Len() > 0 { + out.WriteString("\n") + } + out.WriteString("[stderr]\n") + out.WriteString(stderrText) + } + if err != nil { + if out.Len() > 0 { + out.WriteString("\n") + } + out.WriteString("[运行失败] " + err.Error()) + } + return ToolResult{Success: err == nil, Output: t.truncate(out.String())} +} + +func (t *Toolset) readFile(sessionID, path string) ToolResult { + absPath, err := t.resolvePath(sessionID, path) + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + data, err := os.ReadFile(absPath) + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + return ToolResult{Success: true, Output: t.truncate(string(data))} +} + +func (t *Toolset) writeFile(sessionID, path, content string) ToolResult { + absPath, err := t.resolvePath(sessionID, path) + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + if err := os.WriteFile(absPath, []byte(content), 0o644); err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + return ToolResult{Success: true, Output: fmt.Sprintf("已写入 %d 字节到 %s", len(content), absPath)} +} + +func (t *Toolset) listDirectory(sessionID, path string) ToolResult { + absPath, err := t.resolvePath(sessionID, path) + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + entries, err := os.ReadDir(absPath) + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + sort.Slice(entries, func(i, j int) bool { + return entries[i].Name() < entries[j].Name() + }) + + var out strings.Builder + for _, entry := range entries { + info, infoErr := entry.Info() + if infoErr != nil { + continue + } + name := entry.Name() + if entry.IsDir() { + name += "/" + } + fmt.Fprintf(&out, "%-40s %10d %s\n", name, info.Size(), info.ModTime().Format("2006-01-02 15:04")) + } + if out.Len() == 0 { + return ToolResult{Success: true, Output: "目录为空"} + } + return ToolResult{Success: true, Output: t.truncate(out.String())} +} + +// resolvePath 把工具传入的路径解析为绝对路径: +// - 相对路径基于该会话的独立工作目录(data/workspaces/); +// - 绝对路径允许在项目根目录内(协作黑板产物、agent 自身文件等共享内容); +// - 项目根目录之外一律拒绝。 +func (t *Toolset) resolvePath(sessionID, raw string) (string, error) { + if strings.TrimSpace(raw) == "" { + raw = "." + } + path := filepath.FromSlash(raw) + if !filepath.IsAbs(path) { + path = filepath.Join(t.workspaceFor(sessionID), path) + } + absPath, err := filepath.Abs(path) + if err != nil { + return "", err + } + if !pathWithin(absPath, t.Workspace) { + return "", errors.New("路径超出项目根目录范围,已拒绝: " + absPath) + } + return absPath, nil +} + +// workspaceFor 返回某会话的独立工作目录并确保其存在。 +func (t *Toolset) workspaceFor(sessionID string) string { + root := t.sessionRoot + if root == "" { + root = t.Workspace + } + dir := filepath.Join(root, sessionID) + _ = os.MkdirAll(dir, 0o755) + return dir +} + +// pathWithin 判断 child 是否位于 parent 目录内(含自身)。 +func pathWithin(child, parent string) bool { + rel, err := filepath.Rel(parent, child) + if err != nil { + return false + } + return rel == "." || (rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))) +} + +func (t *Toolset) truncate(value string) string { + runes := []rune(value) + if len(runes) <= t.MaxOutput { + return value + } + return string(runes[:t.MaxOutput]) + "\n...[输出过长,已截断]" +} + +func intParam(params map[string]any, key string) int { + switch value := params[key].(type) { + case float64: + return int(value) + case int: + return value + default: + return 0 + } +} diff --git a/internal/agent/tools_coop.go b/internal/agent/tools_coop.go new file mode 100644 index 0000000..b2d54ac --- /dev/null +++ b/internal/agent/tools_coop.go @@ -0,0 +1,556 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "os" + "os/exec" + "path/filepath" + "regexp" + "strings" + "time" + + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/mount" +) + +// 协作镜像名(构建方式见各自目录的 RUN.md)。 +const ( + coopImagePigo = "pigo-coop" // pigo 协作镜像(支持 openai + anthropic 协议) + coopImage = "pi-coop" // pi 协作镜像(支持 openai + anthropic 协议) + coopImageClaude = "claude-coop" // Claude Code 协作镜像(仅支持 anthropic 协议) +) + +// 协作引擎到镜像 / 容器名前缀 / 展示标签 / 运行时目录的映射。 +// 三种引擎共享同一套环境变量契约(MODEL/BASE_URL/API_KEY/PROTOCOL/TASK 等), +// supervisor 各自适配,控制层只需选择镜像并做协议约束 + baseURL 规范化。 +// runtimeDir 是 local 模式下 supervisor.sh / prompts / extensions 所在的相对目录。 +var coopEngines = map[string]struct { + image, namePrefix, label, runtimeDir string +}{ + "pigo": {coopImagePigo, "pigo-coop-", "pigo", "pigo/coop"}, + "pi": {coopImage, "pi-coop-", "pi", "pi-coop"}, + "claude": {coopImageClaude, "claude-coop-", "claude code", "claude code"}, +} + +// codePattern 从任务描述中提取题目编号(如 "c-06" / "a-05" / "e1-01")。 +var codePattern = regexp.MustCompile(`[A-Za-z0-9]{1,8}-\d{1,4}`) + +// runCoop 异步启动 pi 协作任务: +// 校验参数与镜像后立即返回,容器创建 / 运行 / 等待 / 清理全部放到后台 goroutine。 +// 任务结束后由 CoopManager 通知 Agent,主 agent 会自动收到一条 +// "【协作任务完成通知】"消息并汇报结果,无需在前端长时间等待。 +func (t *Toolset) runCoop(ctx context.Context, sessionID string, params map[string]any) ToolResult { + if t.coop == nil { + return ToolResult{Success: false, Output: "协作任务管理器未初始化"} + } + task, _ := params["task"].(string) + task = strings.TrimSpace(task) + if task == "" { + return ToolResult{Success: false, Output: "task 不能为空"} + } + + // 完成协议:约束协作 agent 主动收尾,避免"已提交成功却没写 DONE 被强杀"、 + // "解不出却空耗到超时"两类问题。supervisor 只在黑板根目录出现 DONE 标记时正常退出。 + completionProtocol := ` + +【完成协议(务必严格遵守,决定协作能否正常收尾)】 +- 找到 flag 并提交成功(平台响应 correct=true)后:用 blackboard 工具 action=done,在黑板根目录创建 DONE 标记,内容写入完成摘要(含 flag 值、提交响应、解题路径)。supervisor 检测到 DONE 即正常结束(exit 0);不创建 DONE 会空耗到超时被强杀,协作被视为失败。 +- 若经充分尝试后确认本轮无法解出(目标不可达 / 无漏洞 / 试错过多):同样用 action=done 创建 DONE,内容开头写明「未解出」与已尝试内容,让调度方及时关闭靶机并切换下一题,不要空耗到超时。 +- 已通关题目不要重复提交:平台对已通关题目的后续提交统一返回 correct:false(而非 duplicate),属正常现象、不影响已得分数,不要误判为失败。` + task += completionProtocol + + // 从任务描述解析题目编号(如 "c-06"),供完成通知携带,便于控制层关靶机/切题 + challengeCode := "" + if m := codePattern.FindString(task); m != "" { + challengeCode = m + } + roundMax := intParam(params, "round_max") + if roundMax <= 0 { + // 单 agent 默认 1 轮:一次运行完成全部工作,未完成则失败并重新下发 + roundMax = 1 + } + if roundMax > 30 { + roundMax = 30 + } + timeoutSec := intParam(params, "timeout") + if timeoutSec <= 0 { + // 默认 1800s:600s/900s 对需要写脚本+多步探测的渗透/解题任务偏紧, + // 实测多因单轮超时(exit 143)导致协作失败。 + timeoutSec = 1800 + } + + if t.apiCfg == nil { + return ToolResult{Success: false, Output: "缺少 LLM API 配置,无法注入模型配置"} + } + + // 解析协作引擎:engine=pi(默认)/ pigo / claude。 + // 未显式指定时使用设置页配置的默认引擎(apiCfg.Engine())。 + // 三种镜像共享同一套环境变量契约(MODEL/BASE_URL/API_KEY/PROTOCOL/TASK 等), + // supervisor 各自适配,控制层只需选择镜像并做协议约束 + baseURL 规范化。 + engine, _ := params["engine"].(string) + engine = strings.TrimSpace(strings.ToLower(engine)) + if engine == "" { + engine = t.apiCfg.Engine() + } + eng, ok := coopEngines[engine] + if !ok { + return ToolResult{Success: false, Output: "不支持的 engine \"" + engine + "\",可选值:pigo | pi(默认)| claude"} + } + + model := t.apiCfg.Model() + apiKey := t.apiCfg.APIKey() + provider := t.apiCfg.Provider() + // coopBaseURL 根据 engine 规范化 baseURL: + // pigo 的 anthropicCompatDriver 只追加 /messages,需补 /v1; + // pi/claude 用官方 SDK(自带 /v1/messages),需剥离 /v1 避免双重路径。 + baseURL := coopBaseURL(provider, t.apiCfg.BaseURL(), engine) + if model == "" || apiKey == "" || baseURL == "" { + return ToolResult{Success: false, Output: "LLM API 配置不完整(model / base_url / api_key 缺一不可),请先在设置页配置"} + } + + // Claude Code 仅支持 Anthropic 协议端点(ANTHROPIC_BASE_URL), + // 若配置为 openai 协议则直接拒绝,避免容器启动后才报错。 + if engine == "claude" && provider != ProviderAnthropic { + return ToolResult{Success: false, Output: "claude-coop 仅支持 Anthropic 协议端点,当前 provider 为 " + provider + + "。请改用 engine=pi 或 engine=pigo,或将 API 配置切换为 anthropic 协议(如 DeepSeek 的 /anthropic 端点)。"} + } + + env := []string{ + "MODEL=" + model, + "BASE_URL=" + baseURL, + "API_KEY=" + apiKey, + "PROTOCOL=" + provider, + "TASK=" + task, + fmt.Sprintf("ROUND_MAX=%d", roundMax), + fmt.Sprintf("TIMEOUT=%d", timeoutSec), + } + + // 生成任务 ID(也用作容器名与默认黑板子目录) + taskID := newID() + + // 黑板目录:用户显式指定 blackboard 时使用指定路径; + // 未指定时为该任务分配独立的会话工作区子目录(data/workspaces//blackboard/), + // 保证同一会话发起的多个 coop 任务互相隔离、互不污染。 + blackboardDir := "" + if dirParam, _ := params["blackboard"].(string); strings.TrimSpace(dirParam) != "" { + dirParam = strings.TrimSpace(dirParam) + abs, err := filepath.Abs(dirParam) + if err != nil { + return ToolResult{Success: false, Output: "blackboard 路径无效: " + err.Error()} + } + // 安全约束:blackboard 会被 worker 进程/容器读写,必须限制在项目根目录内, + // 否则 LLM 可通过指定任意主机目录(如 ~/.ssh)让 worker 读写敏感文件。 + if !pathWithin(abs, t.Workspace) { + return ToolResult{Success: false, Output: "blackboard 路径超出项目根目录范围,已拒绝: " + abs} + } + if err := os.MkdirAll(abs, 0o755); err != nil { + return ToolResult{Success: false, Output: "创建 blackboard 目录失败: " + err.Error()} + } + blackboardDir = abs + } else { + dir := filepath.Join(t.workspaceFor(sessionID), "blackboard", taskID) + if err := os.MkdirAll(dir, 0o755); err != nil { + return ToolResult{Success: false, Output: "创建 blackboard 目录失败: " + err.Error()} + } + blackboardDir = dir + } + + record := &CoopTask{ + ID: taskID, + SessionID: sessionID, + Status: "running", + Blackboard: blackboardDir, + RoundMax: roundMax, + CreatedAt: time.Now(), + ChallengeCode: challengeCode, + } + t.coop.Register(record) + + mode := t.apiCfg.CoopMode() + log.Printf("[coop] 协作任务启动 task=%s session=%s mode=%s engine=%s challenge=%s round_max=%d timeout=%ds blackboard=%s", + record.ID, sessionID, mode, engine, challengeCode, roundMax, timeoutSec, blackboardDir) + + overallSec := timeoutSec*roundMax + 300 + if mode == CoopModeLocal { + // 本地子进程模式:无需 Docker,直接 exec supervisor.sh + supervisor, promptsDir, extensionsDir, err := findCoopRuntime(engine) + if err != nil { + t.coop.Complete(t.failedTask(record, err)) + return ToolResult{Success: false, Output: err.Error()} + } + go t.runCoopLocalAsync(record, env, blackboardDir, supervisor, promptsDir, extensionsDir, roundMax, timeoutSec, eng.label) + } else { + // Docker 容器模式:现有逻辑 + var mounts []mount.Mount + cli, err := t.dockerClient() + if err != nil { + t.coop.Complete(t.failedTask(record, err)) + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + // 判断目标 daemon 系统类型:远程 Linux daemon(如 WSL)的 bind 挂载 + // 只认容器侧路径,Windows 路径需转换为 /mnt/<盘符>/... 形式。 + info, infoErr := cli.Info(ctx) + linuxDaemon := infoErr == nil && info.OSType == "linux" + + // 检查镜像是否存在(同步快速失败,避免后台任务因镜像缺失白跑) + checkCtx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + if _, err := cli.ImageInspect(checkCtx, eng.image); err != nil { + t.coop.Complete(t.failedTask(record, fmt.Errorf("未找到镜像 %s", eng.image))) + return ToolResult{Success: false, Output: "未找到镜像 " + eng.image + "。请先在仓库根构建:\n" + + coopBuildHint(engine) + "\n(详见对应目录的 RUN.md)"} + } + + // 容器内以非 root 用户(uid=1000 agent)运行,主机目录必须对任何用户可写, + // 否则 supervisor 初始化(mkdir/写 task.md/复制 AGENTS.md)会失败。 + _ = os.Chmod(blackboardDir, 0o777) + source := blackboardDir + if linuxDaemon { + source = wslBindPath(blackboardDir) + } + mounts = append(mounts, mount.Mount{ + Type: mount.TypeBind, + Source: source, + Target: "/blackboard", + }) + + // 容器创建 / 启动 / 等待 / 清理放到后台,并使用独立上下文, + // 避免阻塞当前 SSE 流(此前同步等待最长可达 timeout×round_max+300 秒)。 + go t.runCoopAsync(record, env, mounts, roundMax, timeoutSec, eng.image, eng.namePrefix) + } + + return ToolResult{Success: true, Output: fmt.Sprintf( + "协作任务已在后台启动,任务 ID:%s。\n运行方式:%s 协作单 Agent(%s 模式),最多 %d 轮,整体上限约 %d 分钟。\n"+ + "你无需在此等待,可以继续处理其他请求;任务完成后系统会自动通知你并汇报结果。", + record.ID, eng.label, mode, roundMax, overallSec/60)} +} + +// runCoopAsync 在后台完成协作容器的创建、启动、等待、日志收集与清理。 +// 结束(成功 / 失败 / 超时)后通过 coop.Complete 通知 Agent 唤醒主 agent。 +// image / namePrefix 由 runCoop 根据 engine 选择(pi-coop 或 claude-coop)。 +func (t *Toolset) runCoopAsync(record *CoopTask, env []string, mounts []mount.Mount, roundMax, timeoutSec int, image, namePrefix string) { + // 使用独立后台上下文,避免随 SSE 请求断开而中断协作任务 + ctx := context.Background() + + cli, err := t.dockerClient() + if err != nil { + t.coop.Complete(t.failedTask(record, err)) + return + } + defer cli.Close() + + name := namePrefix + record.ID[:8] + created, err := cli.ContainerCreate(ctx, &container.Config{ + Image: image, + Env: env, + }, &container.HostConfig{ + Mounts: mounts, + }, nil, nil, name) + if err != nil { + t.coop.Complete(t.failedTask(record, fmt.Errorf("创建容器失败: %w", err))) + return + } + record.ContainerID = created.ID + + // 运行结束后无论如何清理容器(等价 docker run --rm) + cleanup := func() { + cleanupCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = cli.ContainerRemove(cleanupCtx, created.ID, container.RemoveOptions{Force: true}) + } + + if err := cli.ContainerStart(ctx, created.ID, container.StartOptions{}); err != nil { + cleanup() + t.coop.Complete(t.failedTask(record, fmt.Errorf("启动容器失败: %w", err))) + return + } + + // 等待容器退出;整体超时上限 = 轮次 × 单轮超时 + 缓冲 + overall := time.Duration(timeoutSec*roundMax+300) * time.Second + waitCtx, waitCancel := context.WithTimeout(ctx, overall) + defer waitCancel() + waitCh, errCh := cli.ContainerWait(waitCtx, created.ID, container.WaitConditionNotRunning) + + exitCode := -1 + select { + case res := <-waitCh: + exitCode = int(res.StatusCode) + case err := <-errCh: + cleanup() + t.coop.Complete(t.failedTask(record, fmt.Errorf("等待容器退出失败: %w", err))) + return + case <-waitCtx.Done(): + cleanup() + t.coop.Complete(t.failedTask(record, fmt.Errorf( + "协作运行超时(整体上限 %d 秒,约 %.0f 分钟),已强制清理容器。可增大 round_max / timeout 或缩小任务规模后重试。", + int(overall.Seconds()), overall.Minutes()))) + return + } + + // 读取容器日志(supervisor 输出 + DONE 总结) + logCtx, logCancel := context.WithTimeout(context.Background(), 30*time.Second) + defer logCancel() + result := "" + if logs, err := cli.ContainerLogs(logCtx, created.ID, container.LogsOptions{ + ShowStdout: true, + ShowStderr: true, + }); err == nil { + if raw, readErr := io.ReadAll(logs); readErr == nil { + result = demuxDockerLogs(raw) + _ = logs.Close() + } + } + + result = strings.TrimSpace(result) + switch exitCode { + case 0: + result += "\n[协作完成,退出码 0]" + case 1: + result += "\n[达到最大轮次仍未完成,退出码 1;可增大 round_max 后重试]" + default: + result += fmt.Sprintf("\n[容器退出码 %d,请检查上方日志]", exitCode) + } + if len(mounts) > 0 { + result += "\n黑板产物已保留在: " + mounts[0].Source + } + cleanup() + + record.ExitCode = exitCode + record.Output = t.truncate(result) + + // 解析 supervisor 生成的结构化结果 result.json(如有): + // 完成通知与汇报轮优先使用 Result 字段,容器 stdout 仅作兜底。 + status := "unknown" + if data, rerr := os.ReadFile(filepath.Join(record.Blackboard, "result.json")); rerr == nil { + var res CoopResult + if json.Unmarshal(data, &res) == nil { + record.Result = &res + status = res.Status + } + } + log.Printf("[coop] 协作任务结束 task=%s session=%s exit_code=%d status=%s", record.ID, record.SessionID, exitCode, status) + t.coop.Complete(record) +} + +// failedTask 生成一个失败任务记录,供 runCoopAsync 失败路径统一上报。 +func (t *Toolset) failedTask(record *CoopTask, err error) *CoopTask { + record.Error = err.Error() + record.Output = "协作任务失败: " + err.Error() + return record +} + +// coopBuildHint 返回指定 engine 对应镜像的构建命令提示。 +func coopBuildHint(engine string) string { + switch engine { + case "pigo": + // pigo-coop 构建上下文是 pigo/ 目录,且需先交叉编译 pigo 二进制 + return "cd pigo && GOOS=linux CGO_ENABLED=0 go build -trimpath -ldflags=\"-s -w\" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo && docker build -f coop/Dockerfile -t pigo-coop ." + case "claude": + return "docker build -f \"claude code/Dockerfile\" -t claude-coop \"claude code/\"" + default: + return "docker build -f pi-coop/Dockerfile -t pi-coop pi-coop/" + } +} + +// coopBaseURL 把当前 Agent 配置的 Base URL 规范化为协作容器期望的"基础地址"。 +// 不同 engine 的模型接入层对 baseURL 的路径处理不同,需按 engine 区分: +// - pigo:anthropicCompatDriver 只追加 /messages(不含 /v1),因此 anthropic +// 端点需补 /v1(最终 /v1/messages);openai 端点需保留 /v1。 +// - pi:provider 使用官方 SDK(@anthropic-ai/sdk / openai),SDK 自行追加完整 +// 路径(anthropic: /v1/messages;openai: /chat/completions),不能再补 /v1, +// 否则产生 /v1/v1/messages 双重路径导致 404。 +// - claude:同 pi,用 @anthropic-ai/sdk,不能补 /v1。 +// +// 对用户误填的完整路径后缀(/v1/messages、/chat/completions 等)统一剥离, +// 再按 engine + provider 决定是否补 /v1。 +func coopBaseURL(provider, baseURL, engine string) string { + base := strings.TrimRight(strings.TrimSpace(baseURL), "/") + // 先剥离用户误填的完整路径后缀,统一回退到"基础地址" + switch { + case strings.HasSuffix(base, "/v1/chat/completions"): + base = strings.TrimSuffix(base, "/chat/completions") + case strings.HasSuffix(base, "/chat/completions"): + base = strings.TrimSuffix(base, "/chat/completions") + case strings.HasSuffix(base, "/v1/messages"): + base = strings.TrimSuffix(base, "/v1/messages") + case strings.HasSuffix(base, "/messages"): + base = strings.TrimSuffix(base, "/messages") + } + + if provider != ProviderAnthropic { + return base + } + + // Anthropic 协议端点处理 + switch engine { + case "pigo": + // pigo 的 anthropicCompatDriver 只追加 /messages,需补 /v1。 + // 若用户已填 /v1 结尾则保持,否则补上。 + if strings.HasSuffix(base, "/v1") { + return base + } + return base + "/v1" + default: + // pi / claude 用 @anthropic-ai/sdk,SDK 自带 /v1/messages, + // 不能补 /v1;若用户已填 /v1 则剥离(SDK 会补回完整的 /v1/messages)。 + if strings.HasSuffix(base, "/v1") { + return strings.TrimSuffix(base, "/v1") + } + return base + } +} + +// wslBindPath 把 Windows 绝对路径转换为 WSL 挂载路径(E:\path → /mnt/e/path), +// 供运行在 WSL 内的 Linux Docker daemon 做 bind 挂载。 +func wslBindPath(path string) string { + if len(path) < 2 || path[1] != ':' { + return path + } + drive := strings.ToLower(path[:1]) + rest := strings.ReplaceAll(path[2:], "\\", "/") + return "/mnt/" + drive + rest +} + +// findCoopRuntime 查找指定 engine 的 supervisor.sh / prompts / extensions 目录。 +// 查找顺序:环境变量 COOP_DIR > 相对于可执行文件 > 相对于工作目录。 +// 返回的路径均为绝对路径,供本地子进程模式(runCoopLocalAsync)使用。 +func findCoopRuntime(engine string) (supervisor, prompts, extensions string, err error) { + eng, ok := coopEngines[engine] + if !ok { + return "", "", "", fmt.Errorf("不支持的 engine: %s", engine) + } + dirName := eng.runtimeDir + + // 候选基目录列表:COOP_DIR 环境变量 > 可执行文件同级/上级 > 当前工作目录 + var candidates []string + if envDir := os.Getenv("COOP_DIR"); envDir != "" { + candidates = append(candidates, filepath.Join(envDir, dirName)) + } + if exe, exeErr := os.Executable(); exeErr == nil { + exeDir := filepath.Dir(exe) + candidates = append(candidates, filepath.Join(exeDir, dirName)) + candidates = append(candidates, filepath.Join(exeDir, "..", dirName)) + } + if wd, wdErr := os.Getwd(); wdErr == nil { + candidates = append(candidates, filepath.Join(wd, dirName)) + } + + for _, base := range candidates { + sp := filepath.Join(base, "supervisor.sh") + if st, statErr := os.Stat(sp); statErr == nil && !st.IsDir() { + pp := filepath.Join(base, "prompts") + ep := filepath.Join(base, "extensions") + // prompts / extensions 可选:缺失时传空串,supervisor 用内置默认 + return sp, pp, ep, nil + } + } + return "", "", "", fmt.Errorf( + "本地模式未找到 %s 的 supervisor.sh,已查找目录: %v\n"+ + "请确保 %s 目录存在且包含 supervisor.sh,或设置 COOP_DIR 环境变量指向包含该目录的父目录", + engine, candidates, dirName) +} + +// runCoopLocalAsync 在后台以本地子进程方式运行 supervisor.sh 完成 worker 任务。 +// 与 runCoopAsync(Docker 模式)对应:无需 Docker daemon,直接 exec supervisor.sh, +// 通过 BLACKBOARD/PROMPTS/EXTENSIONS 环境变量指向本地路径。 +// 结束(成功 / 失败 / 超时)后通过 coop.Complete 通知 Agent。 +func (t *Toolset) runCoopLocalAsync(record *CoopTask, env []string, blackboardDir, supervisor, promptsDir, extensionsDir string, roundMax, timeoutSec int, label string) { + // 使用独立后台上下文,避免随 SSE 请求断开而中断协作任务 + ctx := context.Background() + overall := time.Duration(timeoutSec*roundMax+300) * time.Second + runCtx, runCancel := context.WithTimeout(ctx, overall) + defer runCancel() + + // Windows 上 bash 通常是 WSL bash,不认反斜杠路径(E:\foo → E:foo 被吞)。 + // 需把传给 bash 的路径转为 /mnt/<盘符>/... 格式;Go 侧文件操作仍用原始路径。 + toBashPath := func(p string) string { + if len(p) >= 2 && p[1] == ':' { + return wslBindPath(p) + } + return p + } + bashSupervisor := toBashPath(supervisor) + bashBlackboard := toBashPath(blackboardDir) + bashPrompts := toBashPath(promptsDir) + bashExtensions := toBashPath(extensionsDir) + + // 构建子进程环境:继承父进程环境(PATH 等)+ 注入协作环境变量 + procEnv := os.Environ() + procEnv = append(procEnv, env...) + procEnv = append(procEnv, "BLACKBOARD="+bashBlackboard) + if bashPrompts != "" { + procEnv = append(procEnv, "PROMPTS="+bashPrompts) + } + if bashExtensions != "" { + procEnv = append(procEnv, "EXTENSIONS="+bashExtensions) + } + + cmd := exec.CommandContext(runCtx, "bash", bashSupervisor) + cmd.Env = procEnv + // stdout+stderr 合并捕获(supervisor 的日志输出) + var buf bytes.Buffer + cmd.Stdout = &buf + cmd.Stderr = &buf + + log.Printf("[coop] 本地协作进程启动 task=%s pid=pending blackboard=%s supervisor=%s", + record.ID, bashBlackboard, bashSupervisor) + + if err := cmd.Start(); err != nil { + t.coop.Complete(t.failedTask(record, fmt.Errorf("启动 supervisor 失败: %w", err))) + return + } + record.ProcessID = cmd.Process.Pid + log.Printf("[coop] 本地协作进程已启动 task=%s pid=%d", record.ID, record.ProcessID) + + // 等待进程退出(exec.CommandContext 在 runCtx 超时时自动发送 SIGKILL) + waitErr := cmd.Wait() + exitCode := 0 + if waitErr != nil { + if exitErr, ok := waitErr.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else { + exitCode = -1 + } + } + + // 判断是否因超时被杀 + timedOut := runCtx.Err() == context.DeadlineExceeded + + result := strings.TrimSpace(buf.String()) + switch { + case timedOut: + result += fmt.Sprintf("\n[协作运行超时(整体上限 %d 秒,约 %.0f 分钟),已强制终止进程]", + int(overall.Seconds()), overall.Minutes()) + exitCode = 124 + case exitCode == 0: + result += "\n[协作完成,退出码 0]" + case exitCode == 1: + result += "\n[达到最大轮次仍未完成,退出码 1;可增大 round_max 后重试]" + default: + result += fmt.Sprintf("\n[进程退出码 %d,请检查上方日志]", exitCode) + } + result += "\n黑板产物已保留在: " + blackboardDir + + record.ExitCode = exitCode + record.Output = t.truncate(result) + + // 解析 supervisor 生成的结构化结果 result.json(与 docker 模式一致) + status := "unknown" + if data, rerr := os.ReadFile(filepath.Join(blackboardDir, "result.json")); rerr == nil { + var res CoopResult + if json.Unmarshal(data, &res) == nil { + record.Result = &res + status = res.Status + } + } + log.Printf("[coop] 本地协作任务结束 task=%s pid=%d exit_code=%d status=%s", + record.ID, record.ProcessID, exitCode, status) + t.coop.Complete(record) +} diff --git a/internal/agent/tools_docker.go b/internal/agent/tools_docker.go new file mode 100644 index 0000000..79638ce --- /dev/null +++ b/internal/agent/tools_docker.go @@ -0,0 +1,320 @@ +package agent + +import ( + "context" + "encoding/binary" + "encoding/json" + "fmt" + "io" + "sort" + "strconv" + "strings" + + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/image" + "github.com/docker/docker/client" +) + +// Docker 工具:通过 Docker API 操作本地或远程 Docker。 +// socket 地址通过 DockerConfigStore 动态读取,用户在 Web 页面配置后即时生效。 + +func (t *Toolset) dockerClient() (*client.Client, error) { + socket := t.dockerSocket() + cli, err := client.NewClientWithOpts( + client.WithHost(socket), + client.WithAPIVersionNegotiation(), + ) + if err != nil { + return nil, fmt.Errorf("创建 Docker 客户端失败: %w", err) + } + return cli, nil +} + +func (t *Toolset) dockerPS(ctx context.Context, params map[string]any) ToolResult { + all, _ := params["all"].(bool) + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + list, err := cli.ContainerList(ctx, container.ListOptions{All: all}) + if err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + if len(list) == 0 { + return ToolResult{Success: true, Output: "没有找到容器"} + } + + sort.Slice(list, func(i, j int) bool { + return list[i].Created > list[j].Created + }) + var out strings.Builder + fmt.Fprintf(&out, "%-14s %-26s %-20s %-12s %s\n", "CONTAINER ID", "NAME", "IMAGE", "STATUS", "PORTS") + for _, c := range list { + name := strings.TrimPrefix(strings.Join(c.Names, ","), "/") + fmt.Fprintf(&out, "%-14s %-26s %-20s %-12s %s\n", + c.ID[:min(12, len(c.ID))], name, c.Image, c.Status, dockerPorts(c.Ports)) + } + return ToolResult{Success: true, Output: t.truncate(out.String())} +} + +func (t *Toolset) dockerImages(ctx context.Context, params map[string]any) ToolResult { + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + list, err := cli.ImageList(ctx, image.ListOptions{All: true}) + if err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + if len(list) == 0 { + return ToolResult{Success: true, Output: "没有找到镜像"} + } + + sort.Slice(list, func(i, j int) bool { + return list[i].Created > list[j].Created + }) + var out strings.Builder + fmt.Fprintf(&out, "%-22s %-28s %-14s %s\n", "REPOSITORY", "TAG", "IMAGE ID", "SIZE") + for _, img := range list { + if len(img.RepoTags) == 0 { + continue + } + for _, tag := range img.RepoTags { + repo, imageTag, _ := strings.Cut(tag, ":") + id := img.ID + if len(id) > 12 { + id = id[7:19] + } + fmt.Fprintf(&out, "%-22s %-28s %-14s %s\n", repo, imageTag, id, formatBytes(img.Size)) + } + } + return ToolResult{Success: true, Output: t.truncate(out.String())} +} + +func (t *Toolset) dockerLogs(ctx context.Context, params map[string]any) ToolResult { + containerID, _ := params["container"].(string) + containerID = strings.TrimSpace(containerID) + if containerID == "" { + return ToolResult{Success: false, Output: "container 不能为空"} + } + tail := intParam(params, "tail") + if tail <= 0 { + tail = 100 + } + + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + resp, err := cli.ContainerLogs(ctx, containerID, container.LogsOptions{ + ShowStdout: true, + ShowStderr: true, + Tail: strconv.Itoa(tail), + }) + if err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + defer resp.Close() + + raw, err := io.ReadAll(resp) + if err != nil { + return ToolResult{Success: false, Output: "读取 Docker 日志失败: " + err.Error()} + } + return ToolResult{Success: true, Output: t.truncate(demuxDockerLogs(raw))} +} + +// demuxDockerLogs 解析 Docker 日志流:多路复用流(stdout/stderr 分帧)按帧解出, +// 普通文本流(TTY 容器)直接返回。 +func demuxDockerLogs(data []byte) string { + if len(data) >= 8 && isMultiplexedFrame(data[0:8]) { + var out strings.Builder + for i := 0; i+8 <= len(data); { + frameSize := int(binary.BigEndian.Uint32(data[i+4 : i+8])) + i += 8 + if frameSize < 0 || i+frameSize > len(data) { + break + } + out.Write(data[i : i+frameSize]) + i += frameSize + } + return out.String() + } + return string(data) +} + +func isMultiplexedFrame(header []byte) bool { + streamType := header[0] + if streamType != 0 && streamType != 1 && streamType != 2 { + return false + } + return header[1] == 0 && header[2] == 0 && header[3] == 0 +} + +func (t *Toolset) dockerInspect(ctx context.Context, params map[string]any) ToolResult { + containerID, _ := params["container"].(string) + containerID = strings.TrimSpace(containerID) + if containerID == "" { + return ToolResult{Success: false, Output: "container 不能为空"} + } + + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + inspect, err := cli.ContainerInspect(ctx, containerID) + if err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + + state := inspect.State + if state == nil { + state = &container.State{} + } + summary := map[string]any{ + "id": inspect.ID, + "name": strings.TrimPrefix(inspect.Name, "/"), + "image": inspect.Config.Image, + "status": state.Status, + "running": state.Running, + "exit_code": state.ExitCode, + "restart_count": inspect.RestartCount, + "created": inspect.Created, + "command": inspect.Config.Cmd, + "env_count": len(inspect.Config.Env), + "ports": inspect.NetworkSettings.Ports, + "network_mode": string(inspect.HostConfig.NetworkMode), + "ip_address": inspect.NetworkSettings.IPAddress, + } + data, _ := json.MarshalIndent(summary, "", " ") + return ToolResult{Success: true, Output: t.truncate(string(data))} +} + +func (t *Toolset) dockerExec(ctx context.Context, params map[string]any) ToolResult { + containerID, _ := params["container"].(string) + containerID = strings.TrimSpace(containerID) + command, _ := params["command"].(string) + command = strings.TrimSpace(command) + if containerID == "" || command == "" { + return ToolResult{Success: false, Output: "container 和 command 不能为空"} + } + + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + execCfg := container.ExecOptions{ + AttachStdout: true, + AttachStderr: true, + Cmd: []string{"sh", "-c", command}, + } + created, err := cli.ContainerExecCreate(ctx, containerID, execCfg) + if err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + attach, err := cli.ContainerExecAttach(ctx, created.ID, container.ExecAttachOptions{}) + if err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + defer attach.Close() + + raw, err := io.ReadAll(attach.Reader) + if err != nil { + return ToolResult{Success: false, Output: "读取命令输出失败: " + err.Error()} + } + return ToolResult{Success: true, Output: t.truncate(demuxDockerLogs(raw))} +} + +func (t *Toolset) dockerStart(ctx context.Context, params map[string]any) ToolResult { + containerID, _ := params["container"].(string) + containerID = strings.TrimSpace(containerID) + if containerID == "" { + return ToolResult{Success: false, Output: "container 不能为空"} + } + + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + if err := cli.ContainerStart(ctx, containerID, container.StartOptions{}); err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + return ToolResult{Success: true, Output: "容器 " + containerID + " 已启动"} +} + +func (t *Toolset) dockerStop(ctx context.Context, params map[string]any) ToolResult { + containerID, _ := params["container"].(string) + containerID = strings.TrimSpace(containerID) + if containerID == "" { + return ToolResult{Success: false, Output: "container 不能为空"} + } + timeoutSeconds := intParam(params, "timeout") + + cli, err := t.dockerClient() + if err != nil { + return ToolResult{Success: false, Output: err.Error()} + } + defer cli.Close() + + ctx, cancel := context.WithTimeout(ctx, t.Timeout) + defer cancel() + var timeout *int + if timeoutSeconds > 0 { + timeout = &timeoutSeconds + } + if err := cli.ContainerStop(ctx, containerID, container.StopOptions{Timeout: timeout}); err != nil { + return ToolResult{Success: false, Output: "Docker 操作失败: " + err.Error()} + } + return ToolResult{Success: true, Output: "容器 " + containerID + " 已停止"} +} + +func dockerPorts(ports []container.Port) string { + if len(ports) == 0 { + return "" + } + var parts []string + for _, p := range ports { + host := "" + if p.PublicPort > 0 { + host = fmt.Sprintf("%d->", p.PublicPort) + } + parts = append(parts, fmt.Sprintf("%s%d/%s", host, p.PrivatePort, p.Type)) + } + return strings.Join(parts, ", ") +} + +func formatBytes(size int64) string { + const unit = 1024 + if size < unit { + return fmt.Sprintf("%dB", size) + } + div, exp := int64(unit), 0 + for n := size / unit; n >= unit; n /= unit { + div *= unit + exp++ + } + return fmt.Sprintf("%.1f%cB", float64(size)/float64(div), "KMGTPE"[exp]) +} diff --git a/internal/agent/tools_test.go b/internal/agent/tools_test.go new file mode 100644 index 0000000..75975eb --- /dev/null +++ b/internal/agent/tools_test.go @@ -0,0 +1,76 @@ +package agent + +import ( + "path/filepath" + "runtime" + "testing" +) + +// TestPathWithin 验证路径越界校验:read_file / write_file / run_coop 的 blackboard +// 都依赖该函数拦截指向项目根目录之外的路径,是核心安全边界。 +func TestPathWithin(t *testing.T) { + var parent string + if runtime.GOOS == "windows" { + parent = `e:\proj` + } else { + parent = "/proj" + } + + cases := []struct { + name string + child string + expect bool + }{ + {"自身", parent, true}, + {"直接子文件", filepath.Join(parent, "a.txt"), true}, + {"嵌套子目录", filepath.Join(parent, "sub", "deep", "f.txt"), true}, + {"上级目录", filepath.Join(parent, "..", "secret"), false}, + {"同级兄弟目录", filepath.Join(filepath.Dir(parent), "other"), false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := pathWithin(tc.child, parent); got != tc.expect { + t.Fatalf("pathWithin(%q, %q) = %v, want %v", tc.child, parent, got, tc.expect) + } + }) + } +} + +// TestMergeToolCall 验证 SSE 流式 tool_call 增量合并: +// 跨多个 chunk 按 index 聚合 id / name / arguments,是 LLM 工具调用协议正确性的关键。 +func TestMergeToolCall(t *testing.T) { + var calls []ToolCall + // 第一个 chunk:声明 index 0 的调用(id + name) + calls = mergeToolCall(calls, ToolCallDelta{Index: 0, ID: "call_1", Name: "run_bash"}) + // 同一调用的参数分片到达 + calls = mergeToolCall(calls, ToolCallDelta{Index: 0, ArgumentsDelta: `{"comm`}) + calls = mergeToolCall(calls, ToolCallDelta{Index: 0, ArgumentsDelta: `and":"ls"}`}) + + if len(calls) != 1 { + t.Fatalf("应合并为 1 个调用,got %d", len(calls)) + } + c := calls[0] + if c.ID != "call_1" || c.Function.Name != "run_bash" { + t.Fatalf("id/name 不匹配: %+v", c) + } + if c.Function.Arguments != `{"command":"ls"}` { + t.Fatalf("arguments 拼接错误: %q", c.Function.Arguments) + } + + // 第二个调用在 index 1 + calls = mergeToolCall(calls, ToolCallDelta{Index: 1, ID: "call_2", Name: "read_file"}) + calls = mergeToolCall(calls, ToolCallDelta{Index: 1, ArgumentsDelta: `{"path":"a"}`}) + if len(calls) != 2 { + t.Fatalf("应有 2 个调用,got %d", len(calls)) + } + if calls[1].Function.Name != "read_file" { + t.Fatalf("第二个调用名错误: %+v", calls[1]) + } + + // index 为负时回退为追加新调用 + prev := len(calls) + calls = mergeToolCall(calls, ToolCallDelta{Index: -1, ID: "call_3", Name: "list_directory"}) + if len(calls) != prev+1 { + t.Fatalf("负 index 应追加新调用,got %d", len(calls)) + } +} diff --git a/internal/agent/types.go b/internal/agent/types.go new file mode 100644 index 0000000..d2c9aff --- /dev/null +++ b/internal/agent/types.go @@ -0,0 +1,60 @@ +package agent + +import "time" + +type Message struct { + Role string `json:"role"` + Content *string `json:"content"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` +} + +type ToolCall struct { + ID string `json:"id"` + Type string `json:"type"` + Function ToolFunctionCall `json:"function"` +} + +type ToolFunctionCall struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +type Session struct { + ID string `json:"id"` + Title string `json:"title"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + Messages []Message `json:"messages"` +} + +type SessionSummary struct { + ID string `json:"id"` + Title string `json:"title"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + MessageCount int `json:"message_count"` + Preview string `json:"preview"` +} + +type ToolDefinition struct { + Type string `json:"type"` + Function ToolDefinitionFunction `json:"function"` +} + +type ToolDefinitionFunction struct { + Name string `json:"name"` + Description string `json:"description"` + Parameters map[string]any `json:"parameters"` +} + +type ToolResult struct { + Success bool `json:"success"` + Output string `json:"output"` +} + +type Usage struct { + PromptTokens int `json:"prompt_tokens,omitempty"` + CompletionTokens int `json:"completion_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` +} diff --git a/internal/server/auth.go b/internal/server/auth.go new file mode 100644 index 0000000..bab09b2 --- /dev/null +++ b/internal/server/auth.go @@ -0,0 +1,80 @@ +package server + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "net/http" + "strings" + + "github.com/gin-gonic/gin" +) + +// authCookieName 保存已通过授权校验的凭证。值为授权码的 SHA-256 十六进制摘要, +// 不落明文;HttpOnly 使前端 JS 无法读取,降低泄露面。 +const authCookieName = "blackbean_auth" + +// authCookieValue 计算授权码的稳定凭证值(SHA-256 摘要)。 +func authCookieValue(code string) string { + sum := sha256.Sum256([]byte(code)) + return hex.EncodeToString(sum[:]) +} + +// authRequired 是访问授权中间件:cfg.AuthCode 为空(未启用授权)时直接放行; +// 否则要求请求携带与授权码匹配的 Cookie,不匹配则 401 并中止后续处理。 +func (s *Server) authRequired(c *gin.Context) { + if s.cfg.AuthCode == "" { + c.Next() + return + } + if cookie, err := c.Cookie(authCookieName); err == nil && cookie == authCookieValue(s.cfg.AuthCode) { + c.Next() + return + } + fail(c, http.StatusUnauthorized, "需要授权码才能访问,请先在首页输入授权码") + c.Abort() +} + +// authStatus 返回授权状态:required=是否启用了授权,authorized=当前请求是否已通过授权。 +// 该接口始终放行(不挂授权中间件),前端据此决定是否展示授权界面。 +func (s *Server) authStatus(c *gin.Context) { + required := s.cfg.AuthCode != "" + authorized := false + if required { + if cookie, err := c.Cookie(authCookieName); err == nil { + authorized = cookie == authCookieValue(s.cfg.AuthCode) + } + } + ok(c, gin.H{"required": required, "authorized": authorized}) +} + +// authLogin 校验授权码并签发 Cookie。 +// 该接口始终放行(不挂授权中间件);未启用授权时返回 400。 +func (s *Server) authLogin(c *gin.Context) { + if s.cfg.AuthCode == "" { + fail(c, http.StatusBadRequest, "本实例未启用访问授权") + return + } + var req struct { + Code string `json:"code"` + } + if err := c.ShouldBindJSON(&req); err != nil { + fail(c, http.StatusBadRequest, "请求体格式错误") + return + } + // 常量时间比较,避免通过响应时差探测授权码 + want := []byte(s.cfg.AuthCode) + got := []byte(strings.TrimSpace(req.Code)) + if len(want) != len(got) || subtle.ConstantTimeCompare(want, got) != 1 { + fail(c, http.StatusUnauthorized, "授权码错误") + return + } + http.SetCookie(c.Writer, &http.Cookie{ + Name: authCookieName, + Value: authCookieValue(s.cfg.AuthCode), + Path: "/", + HttpOnly: true, + SameSite: http.SameSiteLaxMode, + }) + ok(c, gin.H{"authorized": true}) +} diff --git a/internal/server/benchmark.go b/internal/server/benchmark.go new file mode 100644 index 0000000..415b091 --- /dev/null +++ b/internal/server/benchmark.go @@ -0,0 +1,45 @@ +package server + +import ( + "context" + "log" + + "agent/internal/agent" +) + +// benchmarkInitialTask 是托管模式启动时发给 Agent 的初始任务消息。 +// 系统提示中已包含完整的 Tsecbench 调度规则(并发管理、flag 提交、靶机开关), +// 此处只需一条简短触发指令让 Agent 开始拉取题目并解题。 +const benchmarkInitialTask = `Tsecbench 托管评测已启动。平台已通过环境变量注入 BENCHMARK_TOKEN 和 BENCHMARK_BASE_URL。 + +请立即开始评测: +1. 用 shell 工具读取环境变量:echo $BENCHMARK_BASE_URL && echo $BENCHMARK_TOKEN +2. 用 curl GET $BENCHMARK_BASE_URL/openapi/v1/challenges 获取题目列表 +3. 选定多道未完成题目,逐一 POST $BENCHMARK_BASE_URL/openapi/v1/challenges/start?unique_code=<题号> 启动靶机,把并发拉满(默认上限 3) +4. 对每道题调用 run_coop 下发解题任务(task 中写明目标地址、unique_code、BENCHMARK_TOKEN、提交规则) +5. worker 返回 flag 后用 curl 向平台 submit API 提交,关闭已通关靶机,继续补满并发 +6. 直到所有题目完成或时限到达 + +目标:在时限内尽可能多地完成题目、拿到更高总分。` + +// StartBenchmark 在托管模式下自动创建会话并触发 Agent 开始解题。 +// 由 main.go 在检测到 BENCHMARK_TOKEN 环境变量时以 goroutine 调用。 +// 不依赖 HTTP 服务就绪 —— agent.Run 直接操作 Agent 实例,不经 HTTP 路由。 +// +// 关键:必须注册持久化的 live emitter。coop 完成后 notifyCoopDone 通过 +// liveEmitter 获取 emit 回调来触发汇报轮,若未注册则汇报轮无日志输出、 +// 事件丢失,表现为"评测流程结束"后无任何后续日志(实际仍在运行)。 +// 注册后不注销:整个评测期间(含所有后续汇报轮)保持日志通道畅通。 +func (s *Server) StartBenchmark(ctx context.Context) error { + session := s.store.Create() + sessionID := session.ID + + emit := func(event agent.Event) { + log.Printf("[benchmark] [%s] %v", event.Type, event.Data) + } + // 注册 live emitter,让 notifyCoopDone 触发的汇报轮也能输出日志 + s.agent.RegisterLive(sessionID, emit) + + log.Printf("[benchmark] 自动启动评测会话 %s", sessionID) + return s.agent.Run(ctx, sessionID, benchmarkInitialTask, emit) +} diff --git a/internal/server/handlers_config.go b/internal/server/handlers_config.go new file mode 100644 index 0000000..3b64e70 --- /dev/null +++ b/internal/server/handlers_config.go @@ -0,0 +1,111 @@ +package server + +import ( + "net/http" + "strings" + + "github.com/gin-gonic/gin" +) + +func (s *Server) health(c *gin.Context) { + ok(c, gin.H{"status": "ok"}) +} + +func (s *Server) config(c *gin.Context) { + ok(c, gin.H{ + "model": s.cfg.Model, + "base_url": s.cfg.BaseURL, + "workspace": s.cfg.Workspace, + "max_context_tokens": s.cfg.MaxContextTokens, + "max_tool_result_chars": s.cfg.MaxToolResultChars, + "max_iterations": s.cfg.MaxIterations, + "keep_recent_messages": s.cfg.KeepRecentMessages, + "api_key_configured": s.cfg.APIKey != "", + }) +} + +func (s *Server) dockerConfig(c *gin.Context) { + ok(c, gin.H{ + "docker_socket": s.dockerCfg.Socket(), + "default_socket": s.dockerCfg.DefaultSocket(), + "configured": s.dockerCfg.IsConfigured(), + }) +} + +func (s *Server) setDockerConfig(c *gin.Context) { + var req struct { + Socket string `json:"socket"` + } + if err := c.ShouldBindJSON(&req); err != nil { + fail(c, http.StatusBadRequest, "请求体格式错误") + return + } + req.Socket = strings.TrimSpace(req.Socket) + // 空字符串表示清除配置、恢复默认本地 Docker + if err := s.dockerCfg.SetSocket(req.Socket); err != nil { + fail(c, http.StatusInternalServerError, "保存 Docker 配置失败: "+err.Error()) + return + } + ok(c, gin.H{ + "docker_socket": s.dockerCfg.Socket(), + "default_socket": s.dockerCfg.DefaultSocket(), + "configured": s.dockerCfg.IsConfigured(), + }) +} + +func (s *Server) llmConfig(c *gin.Context) { + apiKey := s.apiCfg.APIKey() + ok(c, gin.H{ + "provider": s.apiCfg.Provider(), + "api_key_configured": apiKey != "", + "api_key_masked": maskSecret(apiKey), + "base_url": s.apiCfg.BaseURL(), + "model": s.apiCfg.Model(), + "engine": s.apiCfg.Engine(), + "coop_mode": s.apiCfg.CoopMode(), + "default_base_url": s.cfg.BaseURL, + "default_model": s.cfg.Model, + }) +} + +func (s *Server) setLLMConfig(c *gin.Context) { + var req struct { + APIKey *string `json:"api_key"` + BaseURL *string `json:"base_url"` + Model *string `json:"model"` + Provider *string `json:"provider"` + Engine *string `json:"engine"` + CoopMode *string `json:"coop_mode"` + } + if err := c.ShouldBindJSON(&req); err != nil { + fail(c, http.StatusBadRequest, "请求体格式错误") + return + } + // 字段为 nil 表示不修改;空字符串表示清除该字段、回退默认 + if err := s.apiCfg.Update(req.APIKey, req.BaseURL, req.Model, req.Provider, req.Engine, req.CoopMode); err != nil { + fail(c, http.StatusInternalServerError, "保存 LLM 配置失败: "+err.Error()) + return + } + apiKey := s.apiCfg.APIKey() + ok(c, gin.H{ + "provider": s.apiCfg.Provider(), + "api_key_configured": apiKey != "", + "api_key_masked": maskSecret(apiKey), + "base_url": s.apiCfg.BaseURL(), + "model": s.apiCfg.Model(), + "engine": s.apiCfg.Engine(), + "coop_mode": s.apiCfg.CoopMode(), + }) +} + +// maskSecret 对密钥做掩码展示,避免在 Web 页面泄露完整值。 +func maskSecret(value string) string { + runes := []rune(value) + if len(runes) <= 6 { + if len(runes) == 0 { + return "" + } + return "****" + } + return string(runes[:3]) + "****" + string(runes[len(runes)-3:]) +} diff --git a/internal/server/handlers_coop.go b/internal/server/handlers_coop.go new file mode 100644 index 0000000..5ff16be --- /dev/null +++ b/internal/server/handlers_coop.go @@ -0,0 +1,63 @@ +package server + +import ( + "strings" + "time" + + "agent/internal/agent" + + "github.com/gin-gonic/gin" +) + +// coopTaskView 是返回给前端的协作任务视图,附带容器的实时状态与协作轮次进度。 +type coopTaskView struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + Status string `json:"status"` // running | done + Error string `json:"error,omitempty"` + ExitCode int `json:"exit_code,omitempty"` + Blackboard string `json:"blackboard,omitempty"` + RoundMax int `json:"round_max"` + CurrentRound int `json:"current_round"` + ContainerState string `json:"container_state,omitempty"` + ContainerStatus string `json:"container_status,omitempty"` + CreatedAt time.Time `json:"created_at"` + FinishedAt time.Time `json:"finished_at,omitempty"` +} + +// coopTasks 返回协作任务列表,支持 ?session= 过滤到某个会话。 +func (s *Server) coopTasks(c *gin.Context) { + tasks := s.agent.CoopTasks(strings.TrimSpace(c.Query("session"))) + views := make([]coopTaskView, 0, len(tasks)) + for _, task := range tasks { + views = append(views, s.coopTaskView(task)) + } + ok(c, gin.H{"tasks": views}) +} + +func (s *Server) coopTaskView(task *agent.CoopTask) coopTaskView { + view := coopTaskView{ + ID: task.ID, + SessionID: task.SessionID, + Status: task.Status, + Error: task.Error, + ExitCode: task.ExitCode, + Blackboard: task.Blackboard, + RoundMax: task.RoundMax, + CreatedAt: task.CreatedAt, + FinishedAt: task.FinishedAt, + } + // 轮次进度以黑板 logs 目录实际出现的轮次为准(容器清理后日志仍在) + view.CurrentRound = agent.CurrentCoopRound(task.Blackboard) + if view.CurrentRound == 0 { + view.CurrentRound = view.RoundMax + } + // 运行中的任务实时查询容器状态;已完成任务容器已被清理,直接读记录即可 + if task.Status == "running" && task.ContainerID != "" { + if state, status, err := agent.ContainerInspectStatus(s.dockerCfg.Socket(), task.ContainerID); err == nil { + view.ContainerState = state + view.ContainerStatus = status + } + } + return view +} diff --git a/internal/server/handlers_sessions.go b/internal/server/handlers_sessions.go new file mode 100644 index 0000000..8d025df --- /dev/null +++ b/internal/server/handlers_sessions.go @@ -0,0 +1,89 @@ +package server + +import ( + "net/http" + "strings" + + "agent/internal/agent" + + "github.com/gin-gonic/gin" +) + +// listSessions 会话列表,支持 ?q=<关键词> 搜索、?limit=&offset= 分页。 +func (s *Server) listSessions(c *gin.Context) { + query := strings.TrimSpace(c.Query("q")) + var sessions []agent.SessionSummary + if query != "" { + sessions = s.store.Search(query) + } else { + sessions = s.store.List() + } + limit, offset := parsePagination(c) + if offset > 0 || limit > 0 { + start := offset + if start > len(sessions) { + start = len(sessions) + } + end := len(sessions) + if limit > 0 && start+limit < end { + end = start + limit + } + sessions = sessions[start:end] + } + ok(c, gin.H{"sessions": sessions}) +} + +func (s *Server) createSession(c *gin.Context) { + session := s.store.Create() + c.JSON(http.StatusCreated, session) +} + +func (s *Server) getSession(c *gin.Context) { + session, found := s.store.Get(c.Param("id")) + if !found { + fail(c, http.StatusNotFound, "会话不存在") + return + } + ok(c, gin.H{"session": session}) +} + +// sessionMessages 分页返回会话消息(offset 从最新一条往前数), +// 响应含 has_more 与 total 供前端"加载更早"。 +func (s *Server) sessionMessages(c *gin.Context) { + id := c.Param("id") + total, exists := s.store.TotalMessages(id) + if !exists { + fail(c, http.StatusNotFound, "会话不存在") + return + } + limit, offset := parsePagination(c) + if limit <= 0 || limit > 200 { + limit = 50 + } + messages, _, hasMore := s.store.Messages(id, limit, offset) + ok(c, gin.H{ + "messages": messages, + "has_more": hasMore, + "total": total, + }) +} + +func (s *Server) deleteSession(c *gin.Context) { + id := c.Param("id") + if !s.store.Delete(id) { + fail(c, http.StatusNotFound, "会话不存在") + return + } + // 释放会话持有的运行时资源(会话锁与 live 注册),避免内存泄漏 + s.agent.ForgetSession(id) + c.Status(http.StatusNoContent) +} + +// clearSession 清空会话历史消息(保留会话)。 +func (s *Server) clearSession(c *gin.Context) { + if !s.store.Clear(c.Param("id")) { + fail(c, http.StatusNotFound, "会话不存在") + return + } + c.Status(http.StatusNoContent) +} diff --git a/internal/server/response.go b/internal/server/response.go new file mode 100644 index 0000000..dd9160f --- /dev/null +++ b/internal/server/response.go @@ -0,0 +1,31 @@ +package server + +import ( + "net/http" + "strconv" + + "github.com/gin-gonic/gin" +) + +// fail 统一错误响应:{ "error": "..." } +func fail(c *gin.Context, status int, message string) { + c.JSON(status, gin.H{"error": message}) +} + +// ok 统一成功响应。 +func ok(c *gin.Context, data gin.H) { + c.JSON(http.StatusOK, data) +} + +// parsePagination 解析 limit / offset 查询参数(非法值归零)。 +func parsePagination(c *gin.Context) (limit, offset int) { + limit, _ = strconv.Atoi(c.Query("limit")) + offset, _ = strconv.Atoi(c.Query("offset")) + if limit < 0 { + limit = 0 + } + if offset < 0 { + offset = 0 + } + return +} diff --git a/internal/server/router.go b/internal/server/router.go new file mode 100644 index 0000000..5a548ae --- /dev/null +++ b/internal/server/router.go @@ -0,0 +1,86 @@ +package server + +import ( + "fmt" + "net/http" + "path/filepath" + "time" + + "agent/internal/agent" + + "github.com/gin-gonic/gin" +) + +// NewRouter 组装全部路由:静态资源、会话直达、/api/v1 REST、/ws WebSocket。 +// 返回 (gin.Engine, *Server):*Server 供 main.go 在托管模式下调用 StartBenchmark 自动触发评测。 +func NewRouter(cfg agent.Config, store *agent.SessionStore, dockerCfg *agent.DockerConfigStore, apiCfg *agent.APIConfigStore) (*gin.Engine, *Server) { + s := &Server{ + cfg: cfg, + store: store, + agent: agent.NewAgent(cfg, store, dockerCfg, apiCfg), + dockerCfg: dockerCfg, + apiCfg: apiCfg, + } + + gin.SetMode(gin.ReleaseMode) + r := gin.New() + r.Use(gin.Recovery(), requestLogger()) + + staticDir := filepath.Join(cfg.Workspace, "web", "static") + indexFile := filepath.Join(staticDir, "index.html") + r.Static("/static", staticDir) + r.StaticFile("/", indexFile) + + // 会话直达:/ 返回前端页面,由前端 JS 从路径解析会话 ID + r.GET("/:sessionId", func(c *gin.Context) { + if !sessionIDPattern.MatchString(c.Param("sessionId")) { + c.Status(http.StatusNotFound) + return + } + c.File(indexFile) + }) + + v1 := r.Group("/api/v1") + v1.GET("/health", s.health) + // 授权相关接口始终放行:前端据此判断是否需要授权 / 提交授权码。 + // 其余 API 与 WebSocket 统一挂 authRequired(未启用授权时自动放行)。 + auth := v1.Group("/auth") + auth.GET("/status", s.authStatus) + auth.POST("/login", s.authLogin) + + authed := v1.Group("") + authed.Use(s.authRequired) + authed.GET("/config", s.config) + authed.GET("/docker/config", s.dockerConfig) + authed.PUT("/docker/config", s.setDockerConfig) + authed.GET("/llm/config", s.llmConfig) + authed.PUT("/llm/config", s.setLLMConfig) + authed.GET("/sessions", s.listSessions) + authed.POST("/sessions", s.createSession) + authed.GET("/sessions/:id", s.getSession) + authed.GET("/sessions/:id/messages", s.sessionMessages) + authed.DELETE("/sessions/:id", s.deleteSession) + authed.DELETE("/sessions/:id/messages", s.clearSession) + authed.GET("/coop/tasks", s.coopTasks) + + // WebSocket 实时通道:Agent 流式输出 / 状态 / 协作通知(同样受授权保护) + r.GET("/ws", s.authRequired, s.ws) + + return r, s +} + +// requestLogger 轻量请求日志中间件:统一 [http] 前缀,便于与 [coop] 等日志区分。 +func requestLogger() gin.HandlerFunc { + return func(c *gin.Context) { + start := time.Now() + c.Next() + gin.DefaultWriter.Write([]byte(fmt.Sprintf( + "[http] %s | %s %s | %s | %d\n", + time.Now().Format("2006/01/02 15:04:05"), + c.Request.Method, + c.Request.URL.Path, + time.Since(start).Round(time.Millisecond).String(), + c.Writer.Status(), + ))) + } +} diff --git a/internal/server/server.go b/internal/server/server.go new file mode 100644 index 0000000..ba5f5da --- /dev/null +++ b/internal/server/server.go @@ -0,0 +1,19 @@ +package server + +import ( + "regexp" + + "agent/internal/agent" +) + +// Server 持有全部依赖,各 handler 文件通过它访问共享状态。 +type Server struct { + cfg agent.Config + store *agent.SessionStore + agent *agent.Agent + dockerCfg *agent.DockerConfigStore + apiCfg *agent.APIConfigStore +} + +// sessionIDPattern 匹配会话直达路径 / WebSocket 绑定的会话 ID 格式(24 位十六进制)。 +var sessionIDPattern = regexp.MustCompile(`^[0-9a-f]{24}$`) diff --git a/internal/server/ws.go b/internal/server/ws.go new file mode 100644 index 0000000..742ff5f --- /dev/null +++ b/internal/server/ws.go @@ -0,0 +1,177 @@ +package server + +import ( + "context" + "encoding/json" + "log" + "net/http" + "strings" + "sync" + "time" + + "agent/internal/agent" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +var upgrader = websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 4096, + // 本地单用户 Web,允许任意 Origin + CheckOrigin: func(r *http.Request) bool { return true }, +} + +// wsInbound 客户端 → 服务端消息。 +// chat:在绑定会话上运行 Agent;switch:切换绑定会话;stop:停止当前生成;ping:保活。 +type wsInbound struct { + Type string `json:"type"` + Content string `json:"content"` + SessionID string `json:"session_id"` +} + +// wsConn 封装一个 WebSocket 连接:绑定一个会话、串行处理消息。 +type wsConn struct { + server *Server + conn *websocket.Conn + sendMu sync.Mutex // 串行化 WriteJSON + + sessionID string + unreg func() // 当前会话的 live 注销函数 + + runMu sync.Mutex + runCancel context.CancelFunc +} + +func (s *Server) ws(c *gin.Context) { + sessionID := strings.TrimSpace(c.Query("session")) + if sessionID == "" || !sessionIDPattern.MatchString(sessionID) { + log.Printf("ws: 连接被拒(session 参数无效)session=%q", sessionID) + fail(c, http.StatusBadRequest, "缺少有效的 session 参数") + return + } + if _, ok := s.store.Get(sessionID); !ok { + log.Printf("ws: 连接被拒(会话不存在)session=%q", sessionID) + fail(c, http.StatusNotFound, "会话不存在") + return + } + + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return + } + client := &wsConn{server: s, conn: conn} + client.bind(sessionID) + client.readPump() +} + +// sendJSON 安全地推送一个事件帧 {type, data}。 +func (cl *wsConn) sendJSON(eventType string, data any) { + cl.sendMu.Lock() + defer cl.sendMu.Unlock() + _ = cl.conn.WriteJSON(map[string]any{"type": eventType, "data": data}) +} + +// bind 切换绑定会话:注销旧会话的 live 注册,注册新会话。 +// 传入空串表示仅注销。 +func (cl *wsConn) bind(sessionID string) { + if cl.unreg != nil { + cl.unreg() + cl.unreg = nil + } + cl.sessionID = sessionID + if sessionID == "" { + return + } + cl.unreg = cl.server.agent.RegisterLive(sessionID, func(event agent.Event) { + cl.sendJSON(event.Type, event.Data) + }) +} + +func (cl *wsConn) handleChat(content string) { + content = strings.TrimSpace(content) + if content == "" { + cl.sendJSON("error", map[string]any{"message": "消息内容不能为空"}) + return + } + // 未配置 LLM API Key 时直接提示去配置,避免进入无效的模型调用。 + // 环境变量 LLM_API_KEY(托管模式)与设置页 api-config.json 都计入已配置。 + if !cl.server.apiCfg.IsAPIKeyConfigured() { + cl.sendJSON("error", map[string]any{"message": "尚未配置 LLM API,请先点击左下角设置按钮配置 API Key 与模型"}) + return + } + cl.runMu.Lock() + if cl.runCancel != nil { + cl.runMu.Unlock() + cl.sendJSON("error", map[string]any{"message": "Agent 正在执行,请等待完成或先停止"}) + return + } + ctx, cancel := context.WithCancel(context.Background()) + cl.runCancel = cancel + cl.runMu.Unlock() + + done := make(chan struct{}) + go func() { + defer close(done) + _ = cl.server.agent.Run(ctx, cl.sessionID, content, func(event agent.Event) { + cl.sendJSON(event.Type, event.Data) + }) + }() + go func() { + <-done + cl.runMu.Lock() + cl.runCancel = nil + cl.runMu.Unlock() + }() +} + +func (cl *wsConn) handleStop() { + cl.runMu.Lock() + if cl.runCancel != nil { + cl.runCancel() + } + cl.runMu.Unlock() +} + +func (cl *wsConn) readPump() { + defer func() { + cl.bind("") // 注销 live + _ = cl.conn.Close() + }() + cl.conn.SetReadLimit(1 << 20) // 1MB + _ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second)) + cl.conn.SetPongHandler(func(string) error { + _ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second)) + return nil + }) + + for { + // 任何收到消息(含客户端心跳 {type:"ping"})都刷新 90s 读超时, + // 否则浏览器只能发协议层 Pong,应用层 ping 不会触发 SetPongHandler,连接必死在第 90 秒 + _ = cl.conn.SetReadDeadline(time.Now().Add(90 * time.Second)) + _, raw, err := cl.conn.ReadMessage() + if err != nil { + return + } + var msg wsInbound + if err := json.Unmarshal(raw, &msg); err != nil { + continue + } + switch msg.Type { + case "chat": + cl.handleChat(msg.Content) + case "stop": + cl.handleStop() + case "switch": + if msg.SessionID != "" && sessionIDPattern.MatchString(msg.SessionID) { + if _, ok := cl.server.store.Get(msg.SessionID); ok { + cl.bind(msg.SessionID) + } + } + case "ping": + cl.sendJSON("pong", map[string]any{"time": time.Now().UnixMilli()}) + default: + log.Printf("ws: unknown message type %q", msg.Type) + } + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..402ba75 --- /dev/null +++ b/main.go @@ -0,0 +1,53 @@ +package main + +import ( + "context" + "log" + "os" + + "agent/internal/agent" + "agent/internal/server" +) + +func main() { + cfg := agent.LoadConfig() + + store, err := agent.NewSessionStore(cfg.DataDir) + if err != nil { + log.Fatalf("failed to initialize session store: %v", err) + } + + dockerCfg, err := agent.NewDockerConfigStore(cfg.DataDir) + if err != nil { + log.Fatalf("failed to initialize docker config store: %v", err) + } + + apiCfg, err := agent.NewAPIConfigStore(cfg.DataDir, agent.APIConfig{ + APIKey: cfg.APIKey, + BaseURL: cfg.BaseURL, + Model: cfg.Model, + Provider: agent.ProviderOpenAI, + }) + if err != nil { + log.Fatalf("failed to initialize api config store: %v", err) + } + + r, srv := server.NewRouter(cfg, store, dockerCfg, apiCfg) + + // 托管模式:检测到 BENCHMARK_TOKEN 时自动启动评测,无需外部触发。 + // agent.Run 不依赖 HTTP 服务,可直接在 goroutine 中调用。 + if os.Getenv("BENCHMARK_TOKEN") != "" { + go func() { + log.Printf("[benchmark] 检测到托管模式(BENCHMARK_TOKEN 已注入),自动启动评测") + if err := srv.StartBenchmark(context.Background()); err != nil { + log.Printf("[benchmark] 初始轮异常: %v", err) + } + log.Printf("[benchmark] 初始轮结束,后续 coop 完成通知将自动触发新一轮处理") + }() + } + + log.Printf("agent is running at http://localhost:%s", cfg.Port) + if err := r.Run(":" + cfg.Port); err != nil { + log.Fatalf("server stopped: %v", err) + } +} diff --git a/pi-coop/AGENTS.md b/pi-coop/AGENTS.md new file mode 100644 index 0000000..fc04760 --- /dev/null +++ b/pi-coop/AGENTS.md @@ -0,0 +1,29 @@ +# 单 agent 任务协议 + +你是独立的任务执行 agent,完成 /blackboard/task.md 中的任务。任务信息通过工作区与黑板目录交换。 + +## 目录结构 + +| 路径 | 用途 | +| --- | --- | +| /blackboard/task.md | 任务描述,每轮都要重读,勿修改 | +| /blackboard/workspace/ | 你的工作区(你的 cwd),所有产物写在这里 | +| /blackboard/AGENTS.md | 本协议(副本在你的工作区,勿修改) | +| /blackboard/DONE | 完成标记,存在即表示任务已交付完成,只用 blackboard 工具创建 | +| /blackboard/logs/ | 每轮运行日志(supervisor 维护) | +| /blackboard/sessions/ | 会话 ID(supervisor 维护,供跨轮恢复) | +| /blackboard/result.json | 任务结果(supervisor 结束时生成,勿手动修改) | + +## 每轮流程 + +1. **读取**:读取 /blackboard/task.md,明确任务与提交规则(若为解题任务,通常含 flag 提交 API / token / unique_code)。 +2. **评估**:检查工作区已有产物,判断进度与缺口,不重复已完成工作。 +3. **执行**:推进任务——执行命令、读写文件、编写产物到 /blackboard/workspace/。 +4. **提交**:拿到 flag 的任务,立即按 task.md 中约定的规则提交,并把提交响应与得分写入产物。 +5. **判定**:全部工作完成、交付物完整后,用 blackboard 工具 action=done 创建 DONE(summary 为最终交付总结,含关键结果/flag/提交响应/产物清单)。DONE 只能创建一次;宁可多轮打磨,不要过早宣布完成;确认无法解出时同样用 done 标记并写明「未解出」与已尝试内容。 + +## 红线 + +- 不直接写 /blackboard/DONE——必须走 blackboard 工具 action=done,保证原子创建。 +- 不伪造实验、命令或提交结果;如实记录。 +- 不修改 /blackboard/task.md、/blackboard/AGENTS.md 与 /blackboard/result.json。 diff --git a/pi-coop/Dockerfile b/pi-coop/Dockerfile new file mode 100644 index 0000000..4753433 --- /dev/null +++ b/pi-coop/Dockerfile @@ -0,0 +1,58 @@ +# pi 协作镜像(单 agent 版,与 blackboard 协议适配) +# +# 构建(在仓库根目录执行,构建上下文为 pi-coop/ 目录): +# cd /mnt/e/Code/Go/awesomeProject/agent +# DOCKER_BUILDKIT=0 docker build -f pi-coop/Dockerfile -t pi-coop pi-coop/ +# +# 基础镜像用 node:22-alpine(自带 node 22 + npm,满足 pi >=22.19 要求)。 +# 在其上装系统工具(curl/wget/git/jq/openssl/ripgrep)和 python3 + 常用库, +# 再全局安装 pi coding agent,最后放入协作组件(extension/supervisor/prompts)。 + +FROM node:22-alpine + +# ---- 系统工具 + Python ---- +# 先把 apk 源换成清华镜像并用 http(node:22-alpine 的 Alpine 3.24 用官方 CDN +# 拉 APKINDEX 会报 TLS: unspecified error,改 http 绕过;装上 ca-certificates 后 +# TLS 即恢复正常)。之后 apk add 装系统工具 + python3。 +RUN sed -i 's|https://dl-cdn.alpinelinux.org|http://mirrors.tuna.tsinghua.edu.cn|g' /etc/apk/repositories \ + && apk add --no-cache \ + bash ca-certificates git ripgrep \ + curl wget jq openssl file \ + python3 py3-pip + +# ---- Python 常用库 ---- +# 用清华 PyPI 镜像加速;--no-cache-dir 避免膨胀;--break-system-packages 放行 Alpine pip 全局安装。 +RUN pip3 install --no-cache-dir --break-system-packages \ + -i https://pypi.tuna.tsinghua.edu.cn/simple \ + requests urllib3 certifi idna charset-normalizer \ + cryptography pyOpenSSL paramiko pyjwt \ + PyYAML beautifulsoup4 lxml \ + python-dotenv click tqdm \ + numpy pandas \ + httpx aiohttp websockets dnspython \ + psutil pillow + +# ---- 全局安装 pi coding agent ---- +# --ignore-scripts 跳过 lifecycle script(pi 官方推荐,dist 已预编译)。 +# 使用 npmmirror 镜像加速国内下载。 +RUN npm install -g --ignore-scripts \ + --registry=https://registry.npmmirror.com \ + @earendil-works/pi-coding-agent@latest + +# ---- 协作组件 ---- +# extension 放 /extensions,supervisor 与 prompts 放默认路径。 +RUN mkdir -p /extensions /prompts /blackboard +COPY coop.ts /extensions/coop.ts +COPY supervisor.sh /usr/local/bin/supervisor.sh +COPY AGENTS.md /prompts/AGENTS.md +COPY prompts /prompts +RUN chmod +x /usr/local/bin/supervisor.sh + +# ---- 非 root 用户 ---- +# node:22-alpine 自带 node 用户(uid 1000),直接复用,与 pigo-coop 的 uid 1000 对齐, +# 便于 blackboard 挂载卷权限一致。把协作目录所有权交给 node。 +RUN chown -R node:node /extensions /prompts /blackboard +USER node +WORKDIR /blackboard + +ENTRYPOINT ["/usr/local/bin/supervisor.sh"] diff --git a/pi-coop/RUN.md b/pi-coop/RUN.md new file mode 100644 index 0000000..ae1f6fd --- /dev/null +++ b/pi-coop/RUN.md @@ -0,0 +1,127 @@ +# pi 协作运行指南 + +> 只讲怎么跑。原理、架构见 [Dockerfile](./Dockerfile) 与 [supervisor.sh](./supervisor.sh)。 + +--- + +## 1. 前置检查(1 分钟) + +```bash +docker --version # 确认 Docker 可用 +node --version # 本机构建不需要,仅容器内 pi 需要 node>=22.19 +``` + +与 pigo-coop 不同,pi-coop **不需要**预编译二进制:pi 通过 npm 全局安装在镜像内, +extension(coop.ts)由 pi 的 jiti 直接加载 .ts,无需预编译。 + +--- + +## 2. 构建 + +在仓库根目录执行(构建上下文为 `pi-coop/` 目录,避免把 `pi/`、`pigo/` 等大目录发给 daemon): + +```bash +docker build -f pi-coop/Dockerfile -t pi-coop pi-coop/ +``` + +构建会: +1. 拉取 `node:24-bookworm-slim` 基础镜像; +2. `npm install -g @earendil-works/pi-coding-agent`(pi 本体); +3. 装系统工具(git/ripgrep/python3/jq/openssl…)与 Python 常用库; +4. 拷入 coop extension、supervisor、协议与提示。 + +> 国内构建慢时,Dockerfile 已用清华 PyPI 镜像加速 pip。npm 若慢可配 `--build-arg` 或宿主 npm 镜像。 + +--- + +## 3. 运行 + +### 3.1 直接运行(不保留任何数据) + +```bash +docker run --rm \ + -e MODEL=deepseek-chat \ + -e BASE_URL=https://api.deepseek.com \ + -e API_KEY= \ + -e PROTOCOL=openai \ + -e TASK="请为 XX 编写设计文档并实现原型" \ + -e ROUND_MAX=6 \ + pi-coop +``` + +### 3.2 推荐运行(保留黑板产物 + 跨 run 会话) + +```bash +mkdir -p blackboard + +docker run --rm \ + -v "$PWD/blackboard:/blackboard" \ + -e MODEL=deepseek-chat \ + -e BASE_URL=https://api.deepseek.com \ + -e API_KEY= \ + -e PROTOCOL=openai \ + -e TASK="请为 XX 编写设计文档并实现原型" \ + -e ROUND_MAX=6 \ + pi-coop +``` + +### 3.3 常用参数速查 + +| 参数 | 必填 | 默认 | 用途 | +| --- | --- | --- | --- | +| `MODEL` | ✅ | — | 模型名 | +| `BASE_URL` | ✅ | — | OpenAI/Anthropic 兼容 base-url | +| `API_KEY` | ✅ | — | API key | +| `PROTOCOL` | 否 | `openai` | `openai` / `anthropic` | +| `TASK` | ✅(或用 `task.md`) | — | 任务描述 | +| `ROUND_MAX` | 否 | `10` | 最大轮次 | +| `TIMEOUT` | 否 | `600` | 单轮超时秒数(`0`=不限) | + +> 环境变量与 pigo-coop 完全一致,`run_coop` 工具无需改动即可注入。 + +--- + +## 4. 结果怎么看 + +- **完成**:stdout 打印 `[supervisor] status=solved`,退出码 `0`。 +- **未完成**:达到 `ROUND_MAX`,`status=unsolved`,退出码 `1`。 +- **参数错**:退出码 `2`。 + +保留黑板后查看: + +```bash +ls blackboard/workspace/ # agent 产物 +cat blackboard/DONE # 最终交付总结 +cat blackboard/result.json # 结构化结果(status/exit_code/summary/flag/artifacts) +cat blackboard/logs/round-1.log # 第 1 轮运行日志(pi --mode json 的 NDJSON 流) +``` + +`result.json` 字段固定:`status`、`exit_code`、`summary`、`flag`、`artifacts`, +与 pigo-coop 完全一致,外部调度方(如 agent-web 的 run_coop 工具)解析逻辑无需改动。 + +--- + +## 5. 与 pigo-coop 的差异 + +| 维度 | pigo-coop | pi-coop | +| --- | --- | --- | +| agent | Go 编译的 pigo 二进制 | npm 安装的 pi(TypeScript) | +| 模型接入 | CLI `--base-url`/`--protocol` 直传 | coop extension 读环境变量注册 provider | +| blackboard 工具 | pigo 内置(blackboard_tool.go) | coop extension 注册(coop.ts) | +| 输出格式 | `-o stream-json` | `--mode json`(NDJSON,首行 session header) | +| session 恢复 | `--resume ` | `--session-id --session-dir ` | +| cwd | `-C ` | 进程 cwd(supervisor 先 `cd`) | +| 协议(task.md/DONE/result.json) | 一致 | 一致 | +| 环境变量(MODEL/BASE_URL/...) | 一致 | 一致 | + +--- + +## 6. 常见运行问题 + +| 报错 / 现象 | 处理 | +| --- | --- | +| 退出 2:`必须提供 MODEL / BASE_URL / API_KEY` | 补传环境变量 | +| `DONE already exists` | 正常,上一次运行已留标记(幂等) | +| `MissingSessionCwdError` | 跨轮 resume 时 session 记录的 cwd(/blackboard/workspace)不存在;确保挂载未变 | +| extension 加载失败 | 确认镜像内 `/extensions/coop.ts` 存在;`pi -e` 用 jiti 加载 ts,无需预编译 | +| 想从头重跑 | 删除挂载目录里的 `DONE` 与 `result.json` 后重新 run | diff --git a/pi-coop/coop.ts b/pi-coop/coop.ts new file mode 100644 index 0000000..9dcc15d --- /dev/null +++ b/pi-coop/coop.ts @@ -0,0 +1,262 @@ +/** + * pi-coop extension:让 pi 适配 blackboard 协作协议。 + * + * 做两件事: + * 1. 注册 "coop" provider —— 从环境变量 BASE_URL / API_KEY / PROTOCOL / MODEL + * 构造一个 OpenAI 兼容或 Anthropic 兼容的 provider,使 supervisor 注入的 + * 模型配置(与 pigo-coop 完全相同的环境变量)能直接被 pi 使用。 + * 2. 注册 "blackboard" 工具 —— 复刻 pigo 的 blackboard 工具(read / post / done), + * 让 agent 通过 action=done 原子创建 /blackboard/DONE 标记,supervisor 据此判定 + * 任务完成。AGENTS.md 协议因此无需改动。 + * + * 与 pigo blackboard_tool.go 行为对齐:post 用 O_APPEND 单次写(原子), + * done 用 O_CREAT|O_EXCL('wx',保证唯一创建),路径全部做 traversal 校验。 + */ +import { Type } from "@earendil-works/pi-ai"; +import { defineTool, type ExtensionAPI } from "@earendil-works/pi-coding-agent"; +import { openSync, closeSync, writeFileSync, appendFileSync, readFileSync, readdirSync, statSync, mkdirSync, existsSync } from "node:fs"; +import { join, resolve, sep, isAbsolute, clean, basename } from "node:path"; + +// 单条 post / 单文件 read 上限:放开到 128KB,复杂任务的脚本与日志能整段读写。 +const MAX_POST_BYTES = 128 * 1024; +const MAX_READ_BYTES = 128 * 1024; + +// provider 协议 → pi api 类型映射。PROTOCOL 由 supervisor 从外部注入 +// (取值 openai | anthropic,见 agent api_config.go Provider())。 +function protocolToApi(protocol: string): "openai-completions" | "anthropic-messages" { + return protocol === "anthropic" ? "anthropic-messages" : "openai-completions"; +} + +export default function (pi: ExtensionAPI) { + const baseUrl = process.env.BASE_URL ?? ""; + const apiKey = process.env.API_KEY ?? ""; + const protocol = process.env.PROTOCOL ?? "openai"; + const modelId = process.env.MODEL ?? ""; + const bbRoot = process.env.BB ?? "/blackboard"; + + if (!baseUrl || !apiKey || !modelId) { + // supervisor 已在前置校验里保证这三个变量非空,这里只做兜底: + // 若 provider 注册失败,pi 启动时会报 model 无法解析,比这里静默更好。 + throw new Error(`coop extension: BASE_URL / API_KEY / MODEL 环境变量必须同时提供`); + } + + // 注册 coop provider。models 数组在工厂里动态构造(id 取自 MODEL), + // 这样 --provider coop --model "$MODEL" 即可命中。 + pi.registerProvider("coop", { + baseUrl, + apiKey, + api: protocolToApi(protocol), + models: [ + { + id: modelId, + name: modelId, + reasoning: true, + input: ["text", "image"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + // 完全放开上下文与输出限制,让复杂任务能充分展开思考与多轮工具调用。 + contextWindow: 128000, + maxTokens: 65536, + }, + ], + }); + + // ---- blackboard 工具 ---- + const blackboardTool = defineTool({ + name: "blackboard", + label: "Blackboard", + description: + "Read and write the task blackboard used by the coop runner. This is the ONLY tool that may touch shared blackboard files atomically. " + + "Actions: read (no path: global snapshot of task.md, workspace listing, DONE state; with path like \"workspace/exploit.py\": contents of that one file); " + + "post (atomically append a progress note; file must be a bare .md name under messages/, e.g. \"round-1-a.md\"); " + + "done (atomically create the DONE marker with a final delivery summary, only when the deliverable is truly complete; fails if DONE already exists). " + + "The blackboard root, current round and your name are available in the environment as BB, ROUND, NAME.", + parameters: Type.Object({ + action: Type.String({ description: "read | post | done" }), + path: Type.Optional(Type.String({ description: 'For read: a path under the blackboard root to fetch (e.g. messages/round-1-a.md). Omit for the global snapshot.' })), + file: Type.Optional(Type.String({ description: 'For post: bare file name under messages/, must end in .md (e.g. round-1-a.md).' })), + content: Type.Optional(Type.String({ description: "For post: the message body (max 128 KiB)." })), + summary: Type.Optional(Type.String({ description: "For done: final delivery summary written into DONE." })), + }), + + async execute(_toolCallId, params) { + const action = String(params.action ?? ""); + try { + switch (action) { + case "read": + return read(params.path); + case "post": + return post(params.file, params.content); + case "done": + return done(params.summary); + default: + return err(`blackboard: unknown action "${action}" (want read|post|done)`); + } + } catch (e) { + return err(`blackboard ${action}: ${e instanceof Error ? e.message : String(e)}`); + } + }, + }); + pi.registerTool(blackboardTool); + + // ---- 实现部分 ---- + + // safePath:把相对路径解析到 blackboard 根内,拒绝绝对路径与 .. 越界。 + // 与 pigo blackboard_tool.go safePath 行为一致。 + function safePath(p: string): string { + const trimmed = (p ?? "").trim(); + if (trimmed === "") throw new Error("empty path"); + const c = clean(trimmed); + if (isAbsolute(c)) throw new Error(`path "${p}" must be relative to the blackboard root`); + const rootClean = clean(bbRoot); + const full = join(rootClean, c); + if (full !== rootClean && !full.startsWith(rootClean + sep)) { + throw new Error(`path "${p}" escapes the blackboard root`); + } + return full; + } + + function ok(text: string) { + return { content: [{ type: "text" as const, text }] }; + } + function err(text: string) { + return { content: [{ type: "text" as const, text }], isError: true }; + } + + function truncate(s: string, max: number): string { + return s.length > max ? s.slice(0, max) + "\n..." : s; + } + + // read:无 path → 全局快照(task.md / messages 列表 / workspace 列表 / DONE 状态); + // 有 path → 单文件内容。 + function read(pathParam: unknown) { + const p = pathParam != null ? String(pathParam) : ""; + if (p.trim() !== "") { + return readFile(p); + } + + const b: string[] = []; + // task.md + try { + const data = readFileSync(join(bbRoot, "task.md"), "utf8"); + b.push("# task.md\n" + truncate(data, MAX_READ_BYTES) + "\n"); + } catch { + b.push("# task.md\n\n"); + } + + // messages/ + b.push("\n# messages/ (" + listNames("messages").join(", ") + ")\n"); + for (const name of listListing("messages")) b.push(" - " + name + "\n"); + + // workspace/ + b.push("\n# workspace/ (your workspace)\n"); + for (const name of listListing("workspace")) b.push(" - " + name + "\n"); + + // DONE + let doneText = ""; + try { + doneText = truncate(readFileSync(join(bbRoot, "DONE"), "utf8"), MAX_READ_BYTES); + } catch {} + b.push("\n# DONE\n"); + b.push(doneText ? "EXISTS:\n" + doneText + "\n" : "\n"); + + // environment + const env: string[] = []; + for (const k of ["BB", "ROUND", "NAME"]) { + const v = (process.env[k] ?? "").trim(); + if (v) env.push(` ${k}=${v}`); + } + if (env.length) b.push("\n# environment\n" + env.join("\n") + "\n"); + + return ok(b.join("")); + } + + function readFile(p: string) { + const full = safePath(p); + let info; + try { + info = statSync(full); + } catch (e) { + return err(`blackboard read: ${p}: ${e instanceof Error ? e.message : String(e)}`); + } + if (info.isDirectory()) return err(`blackboard read: ${p} is a directory; only files can be read`); + const data = readFileSync(full, "utf8"); + return ok("# " + p + "\n" + truncate(data, MAX_READ_BYTES)); + } + + // post:原子追加一条消息到 messages/。file 必须是裸 *.md 名。 + function post(fileParam: unknown, contentParam: unknown) { + const name = String(fileParam ?? "").trim(); + if (!validMessageName(name)) { + return err('blackboard post: file must be a bare name ending in .md (e.g. "round-1-a.md"), no path separators, no ".."'); + } + const content = String(contentParam ?? "").trim(); + if (content === "") return err("blackboard post: content must not be empty"); + if (Buffer.byteLength(content) > MAX_POST_BYTES) { + return err(`blackboard post: content too large (${Buffer.byteLength(content)} bytes, max ${MAX_POST_BYTES})`); + } + const dir = join(bbRoot, "messages"); + mkdirSync(dir, { recursive: true }); + // O_APPEND 单次写:POSIX 下并发 post 不会交错字节。 + appendFileSync(join(dir, name), content + "\n"); + return ok(`Message appended to messages/${name}`); + } + + // done:原子创建 DONE 标记。'wx' = O_CREAT|O_EXCL|O_WRONLY,保证唯一创建。 + function done(summaryParam: unknown) { + const summary = String(summaryParam ?? "").trim(); + if (summary === "") return err("blackboard done: summary must not be empty (include the final delivery summary)"); + const header = "Blackboard cooperation DONE\ncreated: " + new Date().toISOString() + "\n\n"; + const target = join(bbRoot, "DONE"); + let fd: number; + try { + fd = openSync(target, "wx"); + } catch (e) { + if (existsSync(target)) { + const existing = readFileSync(target, "utf8"); + return err("blackboard done: DONE already exists — cooperation already finished:\n" + truncate(existing, MAX_READ_BYTES)); + } + return err("blackboard done: " + (e instanceof Error ? e.message : String(e))); + } + try { + writeFileSync(fd, header + summary); + } finally { + closeSync(fd); + } + return ok("DONE marker created. Cooperation finished."); + } + + // 校验 post 文件名:裸 *.md,无分隔符,无 . / .. + function validMessageName(name: string): boolean { + if (name === "" || !name.endsWith(".md")) return false; + if (/[\\/]/.test(name) || name === "." || name === "..") return false; + const base = name.slice(0, -3); + if (base === "" || base.startsWith(".") || base.includes("..")) return false; + return true; + } + + function listNames(sub: string): string[] { + try { + return readdirSync(join(bbRoot, sub)).sort(); + } catch { + return []; + } + } + + function listListing(sub: string): string[] { + try { + return readdirSync(join(bbRoot, sub)) + .sort() + .map((name) => { + try { + const info = statSync(join(bbRoot, sub, name)); + if (info.isDirectory()) return `${name} (dir)`; + return `${name} (${info.size} bytes, ${info.mtime.toISOString().slice(11, 19)})`; + } catch { + return name; + } + }); + } catch { + return []; + } + } +} diff --git a/pi-coop/prompts/agent.md b/pi-coop/prompts/agent.md new file mode 100644 index 0000000..d830e6d --- /dev/null +++ b/pi-coop/prompts/agent.md @@ -0,0 +1,23 @@ +# 角色:agent(主执行者) + +你是 pi 单 agent 任务执行者,独立完成 /blackboard/task.md 中的任务。 + +## 容器内可用工具 +容器内预装了安全工具(nuclei、observer_ward、radare2、gdb、ROPgadget、capstone 等)。 +**解题前先读取 `/opt/tools/TOOLS.md` 了解完整工具清单和用法**,选择最合适的工具。 +也可用 `which ` 或 `pip3 list` 确认工具是否可用。 + +## 你的职责 +- 理解任务、制定方案、完成核心产出。 +- 所有产物写入你的工作区(你的 cwd,即 /blackboard/workspace/),按需命名。 +- 需要提交 flag 的任务:拿到 flag 后立即按任务中的提交规则提交,并把提交响应与得分记录在产物中。 + +## 工作方式 +- 每轮:读取 /blackboard/task.md 与工作区已有产物 → 推进任务 → 将进展与产物写入工作区。 +- 环境变量 ROUND、NAME、BB 由 supervisor 提供;需要记录进展或创建完成标记时用 blackboard 工具(BB 指向黑板根目录)。 +- 不要伪造命令结果、文件内容或提交响应;如实记录。 +- 全部工作完成、交付物完整时,用 blackboard 工具 action=done 创建完成标记(summary 写最终交付总结,含关键结果、flag、提交响应、产物清单)。 + +## 注意 +- 若本轮未完成,下一轮会用 --session-id 恢复你的会话继续推进,跨轮保持上下文。 +- 宁可多轮打磨,不要过早宣布完成;确认无法解出时同样用 action=done 标记,summary 写明「未解出」与已尝试内容,让调度方及时切换下一题。 diff --git a/pi-coop/supervisor.sh b/pi-coop/supervisor.sh new file mode 100644 index 0000000..25b1bb9 --- /dev/null +++ b/pi-coop/supervisor.sh @@ -0,0 +1,207 @@ +#!/usr/bin/env bash +# pi 协作 supervisor(单 agent 版) +# +# 在容器内运行一个 pi headless 进程完成 /blackboard/task.md 中的任务。 +# agent 的工作区为 /blackboard/workspace,产物与轮次日志保留在黑板上; +# 任务完成(agent 用 blackboard 工具创建 DONE)后,supervisor 把结果 +# 以结构化 JSON(result.json)落盘并打印一行摘要,供外部调度方直接解析。 +# +# 与 pigo-coop supervisor 的差异(因 pi CLI 与 pigo 不同): +# 1. cwd:pi 用进程 cwd 作工作区(无 -C 参数),supervisor 先 cd 到 +# $BB/workspace 再启动 pi。 +# 2. 模型接入:pi 用 provider 体系,没有 --base-url/--protocol。 +# 由 coop extension(-e /extensions/coop.ts)读 BASE_URL/API_KEY/ +# PROTOCOL/MODEL 环境变量注册 "coop" provider;CLI 用 +# --provider coop --model "$MODEL"。 +# 3. 输出格式:pi 用 --mode json(NDJSON),首行是 session header +# {"type":"session","id":"",...},supervisor 从首行提取 id +# 供下一轮 --session-id 恢复。 +# 4. session 恢复:pi 用 --session-id + --session-dir , +# 替代 pigo 的 --resume 。 +# +# 环境变量(均由外部调用方传入,与 pigo-coop 完全一致): +# MODEL 模型名,如 deepseek-chat (必填) +# BASE_URL OpenAI/Anthropic 兼容 API base-url (必填) +# API_KEY API key (必填) +# PROTOCOL 协议 openai|anthropic (可选,默认按 model 推断) +# TASK 任务描述 (必填;或挂载 $BB/task.md) +# BLACKBOARD 黑板目录 (默认 /blackboard) +# PROMPTS 协作 prompt 目录 (默认 /prompts) +# EXTENSIONS extension 目录 (默认 /extensions) +# ROUND_MAX 最大轮次 (默认 20) +# TIMEOUT 单轮超时秒数,0=不超时 (默认 1800) +# FAIL_MODE agent 失败时:stop=立即退出 | continue=继续下一轮(默认) +set -u + +BB="${BLACKBOARD:-/blackboard}" +PROMPTS="${PROMPTS:-/prompts}" +EXTENSIONS="${EXTENSIONS:-/extensions}" +ROUND_MAX="${ROUND_MAX:-20}" +TIMEOUT="${TIMEOUT:-1800}" +FAIL_MODE="${FAIL_MODE:-continue}" + +log() { echo "[supervisor] $*"; } + +# emit_result 把任务结果以结构化 JSON 写入 $BB/result.json 并打印一行摘要。 +# status: solved | unsolved | error | timeout +emit_result() { + local status="$1" summary="$2" code="$3" + local flag="" artifacts_json="[]" summary_json='""' flag_json='""' + + # flag 优先从 summary 提取,其次在 workspace 产物中全量检索 + if [ -n "$summary" ]; then + flag=$(printf '%s' "$summary" | grep -oE 'flag\{[^}]+\}' | head -1) + fi + if [ -z "$flag" ] && [ -d "$BB/workspace" ]; then + flag=$(grep -rhoE 'flag\{[^}]+\}' "$BB/workspace" 2>/dev/null | head -1) + fi + + # 产物清单:workspace 下的全部文件(最多 200 个),经 python3 转义为 JSON 数组 + if [ -d "$BB/workspace" ]; then + artifacts_json=$( + cd "$BB/workspace" && find . -type f 2>/dev/null | sed 's|^\./||' | head -200 | + python3 -c 'import json,sys; print(json.dumps([l.rstrip("\n") for l in sys.stdin]))' 2>/dev/null + ) + [ -z "$artifacts_json" ] && artifacts_json="[]" + fi + + # summary 截断到 8192 字符并经 python3 转义,防止引号/换行破坏 JSON + summary_json=$(printf '%s' "$summary" | head -c 8192 | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))' 2>/dev/null) + [ -z "$summary_json" ] && summary_json='""' + flag_json=$(printf '%s' "$flag" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))' 2>/dev/null) + [ -z "$flag_json" ] && flag_json='""' + + cat > "$BB/result.json" <&2 + exit 2 +fi +if [ -z "${TASK:-}" ] && [ ! -f "$BB/task.md" ]; then + echo "错误:请通过 TASK 环境变量或挂载 $BB/task.md 提供任务" >&2 + exit 2 +fi +if [ -f "$BB/DONE" ]; then + log "黑板已有完成标记,如需重新开始请删除 $BB/DONE" + exit 0 +fi + +# ---- 初始化黑板 ---- +mkdir -p "$BB/workspace" "$BB/logs" "$BB/sessions" +if [ -n "${TASK:-}" ]; then + printf '%s\n' "$TASK" > "$BB/task.md" +fi +# 协议注入:agent 的 cwd 是 $BB/workspace,AGENTS.md 注入链从 cwd 起, +# 因此把协议副本放进工作区(勿修改,它是每轮系统提示的一部分) +cp "$PROMPTS/AGENTS.md" "$BB/workspace/AGENTS.md" +log "黑板初始化完成:$BB" +log "任务:$(head -c 200 "$BB/task.md")" + +# 从 pi --mode json 输出的首行(session header)提取 session id。 +# 首行形如 {"type":"session","version":3,"id":"","cwd":...} +extract_session_id() { + local logfile="$1" + head -1 "$logfile" 2>/dev/null | python3 -c ' +import json, sys +try: + d = json.loads(sys.stdin.read()) + if d.get("type") == "session": + print(d.get("id", "")) +except Exception: + pass +' 2>/dev/null +} + +# ---- 单轮运行:一个 pi 进程,--mode json 输出 NDJSON(首行携带 session id)---- +run_agent() { + local round="$1" + local session_file="$BB/sessions/agent.session" + local task_prompt="轮次 $round 开始。先读取 /blackboard/task.md 中的任务,检查工作区已有产物与 DONE 状态,然后继续推进任务。全部工作完成、交付物完整时,用 blackboard 工具 action=done 创建完成标记(summary 写最终交付总结,含关键结果/flag/提交响应/产物清单)。" + local args=(-p "$task_prompt" -a --mode json \ + --provider coop --model "$MODEL" \ + --session-dir "$BB/sessions" \ + -e "$EXTENSIONS/coop.ts" \ + --append-system-prompt "$PROMPTS/agent.md") + if [ -s "$session_file" ]; then + # --session-id:精确匹配,找不到则用该 id 新建(不报错),比 --session 更健壮 + args+=(--session-id "$(cat "$session_file")") + fi + + export ROUND="$round" NAME="agent" BB="$BB" + local logfile="$BB/logs/round-$round.log" + log "第 $round 轮开始运行 agent(日志:$logfile)" + # pi 用进程 cwd 作为工作区,必须先进入 $BB/workspace + if [ "$TIMEOUT" -gt 0 ] 2>/dev/null; then + (cd "$BB/workspace" && timeout "$TIMEOUT" pi "${args[@]}") > "$logfile" 2>&1 + else + (cd "$BB/workspace" && pi "${args[@]}") > "$logfile" 2>&1 + fi + local rc=$? + + # 从 --mode json 首行(session header)提取 session id 供下一轮恢复 + local sid + sid=$(extract_session_id "$logfile") + if [ -n "$sid" ]; then + printf '%s' "$sid" > "$session_file" + fi + return $rc +} + +# ---- 主循环:逐轮运行 agent,直到 DONE 或达到轮次上限 ---- +for round in $(seq 1 "$ROUND_MAX"); do + [ -f "$BB/DONE" ] && break + log "===== 第 $round 轮开始 =====" + + run_agent "$round" + rc=$? + + if [ "$rc" -eq 0 ]; then + log "第 $round 轮完成" + else + log "第 $round 轮失败(exit=$rc),日志见 $BB/logs/round-$round.log" + if [ "$FAIL_MODE" = "stop" ]; then + # 兜底:agent 虽超时/失败,但黑板已有提交成功证据(workspace 内 *.md 含 correct:true) + # → 视为完成正常退出,避免"flag 已提交却因未写 DONE 被强杀(exit 143)"。 + _done=0 + for _f in "$BB"/workspace/*.md "$BB"/messages/*.md; do + [ -f "$_f" ] && grep -q '"correct":true' "$_f" && _done=1 && break + done + if [ "$_done" = "1" ]; then + log "检测到提交成功证据,自动标记完成" + { echo "Task finished (auto-detected submit success)"; } > "$BB/DONE" + emit_result solved "$(cat "$BB/DONE")" 0 + exit 0 + fi + if [ "$rc" -eq 124 ]; then + emit_result timeout "" "$rc" + else + emit_result error "agent 退出码 $rc,日志见 logs/round-$round.log" "$rc" + fi + exit "$rc" + fi + fi +done + +if [ -f "$BB/DONE" ]; then + log "检测到完成标记,任务结束" + emit_result solved "$(cat "$BB/DONE")" 0 + exit 0 +fi + +log "达到最大轮次 $ROUND_MAX 仍未完成,请检查黑板产物与日志" +emit_result unsolved "" 1 +exit 1 diff --git a/pigo/.dockerignore b/pigo/.dockerignore new file mode 100644 index 0000000..8e03c29 --- /dev/null +++ b/pigo/.dockerignore @@ -0,0 +1,10 @@ +# 精简 Docker 构建上下文(用于 coop/Dockerfile) +.git +book +docs +examples +tasks +*.exe +pigo +install.sh +vm_vpnstart.py diff --git a/pigo/.gitignore b/pigo/.gitignore new file mode 100644 index 0000000..c7c6b42 --- /dev/null +++ b/pigo/.gitignore @@ -0,0 +1,62 @@ +# If you prefer the allow list template instead of the deny list, see community template: +# https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore +# +# Binaries for programs and plugins +*.exe +*.exe~ +*.dll +*.so +*.dylib + +# Test binary, built with `go test -c` +*.test + +# Code coverage profiles and other test artifacts +*.out +coverage.* +*.coverprofile +profile.cov + +# Dependency directories (remove the comment below to include it) +# vendor/ + +# Go workspace file +go.work +go.work.sum + +# env file +.env + +# pigo config files (may contain plaintext API keys) +config.toml +**/config.toml +.pigo/ + +# Editor/IDE +# .idea/ +# .vscode/ + +.loop-state.json + +# Compiled binary (go build ./cmd/pigo produces ./pigo) +/pigo + +# goreleaser output +/dist/ +.graph_state +graph.html + +# Book build products (source in, products out — see book/README.md) +book/*.pdf +book/book.tex +book/.build-bin/ +book/*.aux +book/*.log +book/*.toc +book/*.out +book/*.bcf +book/*.bbl +book/*.blg +book/*.run.xml +# Figures pre-rendered from SVG by build_pdf.sh (rsvg-convert); source is the .svg +book/images/*.pdf diff --git a/pigo/.goreleaser.yaml b/pigo/.goreleaser.yaml new file mode 100644 index 0000000..66863a3 --- /dev/null +++ b/pigo/.goreleaser.yaml @@ -0,0 +1,86 @@ +# goreleaser 配置:构建跨平台二进制、打包、生成 checksums、发布到 GitHub Release。 +# +# 本地校验: goreleaser check +# 本地试跑(不发布): goreleaser release --snapshot --clean +# 正式发布: 打 tag 后 `git push --tags`,CI 触发 `goreleaser release --clean` +version: 2 + +project_name: pigo + +before: + hooks: + - go mod tidy + +builds: + - id: pigo + main: ./cmd/pigo + binary: pigo + env: + - CGO_ENABLED=0 + goos: + - linux + - darwin + - windows + goarch: + - amd64 + - arm64 + # version/commit/date 变量注入 main 包(见 cmd/pigo/main.go)。 + ldflags: + - -s -w + - -X main.version={{.Version}} + - -X main.commit={{.Commit}} + - -X main.date={{.Date}} + +archives: + - id: pigo + ids: + - pigo + name_template: >- + {{ .ProjectName }}_{{ .Version }}_ + {{- if eq .Os "darwin" }}Darwin + {{- else if eq .Os "linux" }}Linux + {{- else if eq .Os "windows" }}Windows + {{- else }}{{ .Os }}{{ end }}_ + {{- if eq .Arch "amd64" }}x86_64 + {{- else if eq .Arch "386" }}i386 + {{- else }}{{ .Arch }}{{ end }} + format_overrides: + - goos: windows + formats: + - zip + files: + - README.md + - LICENSE + +checksum: + name_template: "checksums.txt" + +snapshot: + version_template: "{{ incpatch .Version }}-next" + +changelog: + sort: asc + use: github + filters: + exclude: + - "^docs:" + - "^test:" + - "^chore:" + - "Merge pull request" + - "Merge branch" + groups: + - title: Features + regexp: '^.*?feat(\(.+\))??!?:.+$' + order: 0 + - title: Bug fixes + regexp: '^.*?fix(\(.+\))??!?:.+$' + order: 1 + - title: Others + order: 999 + +release: + github: + owner: smallnest + name: pigo + draft: false + prerelease: auto diff --git a/pigo/LICENSE b/pigo/LICENSE new file mode 100644 index 0000000..5aaeebe --- /dev/null +++ b/pigo/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 smallnest + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/pigo/agent/agent.go b/pigo/agent/agent.go new file mode 100644 index 0000000..98f8148 --- /dev/null +++ b/pigo/agent/agent.go @@ -0,0 +1,165 @@ +package agent + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" +) + +// Session is a single, stateful pigo agent conversation. Create one with New, +// drive it with Prompt or Stream, and release its resources with Close. The +// conversation history accumulates across calls, so follow-up prompts see the +// earlier exchange; call Reset to start over on the same session. +// +// A Session is not safe for concurrent use. Drive it from a single goroutine, or +// give each goroutine its own Session. +type Session struct { + env run.Env + runCfg runtime.RunConfig + agentCtx *agentcore.AgentContext + model string +} + +// New builds a Session from the given options. It resolves the provider and +// credentials, assembles the tool set, and validates the tool policy and +// thinking level up front, so a configuration mistake (an unknown tool name, an +// invalid thinking level, an unresolvable provider) is returned here rather than +// surfacing on the first Prompt. +// +// No network call is made by New: the provider is only contacted when you call +// Prompt or Stream. This makes New cheap and safe to use in tests. +// +// See the package documentation for the default tool, skill, and memory +// behavior — in particular, that tools are enabled and auto-executed by default. +func New(opts ...Option) (*Session, error) { + c := config{model: "openrouter/free"} + for _, o := range opts { + o(&c) + } + + // Validate the reasoning-effort level through the same layered config chain + // the CLI uses, so an invalid WithThinkingLevel value fails fast here. + thinking, err := run.ResolveThinkingLevel(c.thinking) + if err != nil { + return nil, err + } + + // One ToolPolicy value carries both lists so they cannot be swapped; deny + // always wins over allow inside run.ApplyToolPolicy. + policy := run.NewToolPolicy(c.allowedTools, c.disallowedTools) + + // SetupEnv resolves the provider, assembles the (policy-filtered) tool set, + // builds the system prompt, and — because skills/memory are opt-in here — + // leaves the machine's shared state untouched unless WithSkills/WithMemory + // were passed. It also validates the tool policy against the real tool set, + // so an unknown tool name is reported as an error. + env, err := run.SetupEnv( + c.model, c.baseURL, c.protocol, c.provider, c.apiKey, + c.noTools, !c.skills, c.systemPrompt, c.appendSystemPrompt, c.memory, policy, + ) + if err != nil { + return nil, err + } + + // Resolve the API key by provider name: an explicit WithAPIKey overrides the + // provider's environment variable. The key is held only in the credential + // store and never logged. + creds := provider.NewCredentialStore(nil) + if c.apiKey != "" { + creds.SetOverride(env.ProviderName, c.apiKey) + } + + runCfg := run.NewConfig( + c.model, env.ProviderName, thinking, env.Provider, creds, + run.ToolRegistry(env.Tools), run.TodoReminders(env.Tools), + ) + + return &Session{ + env: env, + runCfg: runCfg, + agentCtx: &agentcore.AgentContext{ + SystemPrompt: env.SysPrompt, + Tools: env.Tools, + }, + model: c.model, + }, nil +} + +// Prompt sends one user message, runs the agent loop to completion (executing +// any tool calls the model makes along the way), and returns the assistant's +// final text. The exchange is appended to the session history so later prompts +// have this context. +func (s *Session) Prompt(ctx context.Context, prompt string) (string, error) { + return s.Stream(ctx, prompt, nil) +} + +// Stream is Prompt with incremental output: onText, if non-nil, is called with +// each chunk of assistant text as it arrives, and the complete final text is +// also returned. Tool calls still run automatically between text chunks. A nil +// onText makes Stream behave exactly like Prompt. +func (s *Session) Stream(ctx context.Context, prompt string, onText func(string)) (string, error) { + // The loop expects the initiating user message already appended; it then + // mutates agentCtx.Messages in place (assistant + tool results), which is + // what carries the conversation forward across calls. + s.agentCtx.Messages = append(s.agentCtx.Messages, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(prompt)}, + }) + + stream := runtime.StartRun(ctx, s.agentCtx, s.runCfg) + final, err := runtime.DrainStream(ctx, stream, runtime.StreamHandler{OnText: onText}) + if err != nil { + return "", err + } + if final == nil { + return "", nil + } + return agentcore.ContentToText(final.Content), nil +} + +// Reset clears the conversation history, so the next Prompt starts a fresh +// exchange. The provider, tool set, and system prompt are unchanged. +func (s *Session) Reset() { + s.agentCtx.Messages = nil +} + +// ToolNames returns the names of the tools available to this session, in the +// order they are advertised to the model. It reflects the applied tool policy, +// so it is a convenient way to confirm WithTools/WithDisallowedTools did what +// you intended. The result is empty for a WithoutTools session. +func (s *Session) ToolNames() []string { + names := make([]string, len(s.env.Tools)) + for i, t := range s.env.Tools { + names[i] = t.Name() + } + return names +} + +// Model returns the model id the session was created with. +func (s *Session) Model() string { return s.model } + +// Provider returns the resolved provider name (e.g. "anthropic", "openrouter"), +// which is inferred from the model id unless WithProvider was set. +func (s *Session) Provider() string { return s.env.ProviderName } + +// Close releases resources held by the session: any loaded plugin manager and +// the persistent memory store (when WithMemory was used). It is safe to call +// once, and safe to call on a session that holds neither. After Close the +// session must not be used again. +func (s *Session) Close() error { + var firstErr error + if s.env.Plugins != nil { + if err := s.env.Plugins.Close(); err != nil && firstErr == nil { + firstErr = err + } + } + if s.env.Memory != nil { + if err := s.env.Memory.Close(); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} diff --git a/pigo/agent/agent_test.go b/pigo/agent/agent_test.go new file mode 100644 index 0000000..6de5478 --- /dev/null +++ b/pigo/agent/agent_test.go @@ -0,0 +1,160 @@ +package agent_test + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/agent" +) + +// hermetic points provider/skill/plugin discovery at throwaway dirs and supplies +// a dummy key so New resolves fully without ever contacting a network. None of +// these tests call Prompt/Stream, so no real request is made. +func hermetic(t *testing.T) { + t.Helper() + t.Setenv("OPENROUTER_API_KEY", "test-key") + t.Setenv("PIGO_HOME", t.TempDir()) +} + +func contains(set []string, name string) bool { + for _, n := range set { + if n == name { + return true + } + } + return false +} + +// TestNewDefaults is the zero-config path: the full built-in tool set is +// advertised and the model id resolves to the openrouter provider. +func TestNewDefaults(t *testing.T) { + hermetic(t) + sess, err := agent.New(agent.WithModel("openrouter/free")) + if err != nil { + t.Fatalf("New: %v", err) + } + defer sess.Close() + + if got := sess.Model(); got != "openrouter/free" { + t.Errorf("Model() = %q, want %q", got, "openrouter/free") + } + if got := sess.Provider(); got != "openrouter" { + t.Errorf("Provider() = %q, want %q", got, "openrouter") + } + for _, want := range []string{"read", "write", "edit", "grep", "find", "bash", "task"} { + if !contains(sess.ToolNames(), want) { + t.Errorf("default tool set missing %q: %q", want, sess.ToolNames()) + } + } +} + +// TestWithToolsAllowlist confirms an allowlist narrows the set to exactly the +// named tools, in order. +func TestWithToolsAllowlist(t *testing.T) { + hermetic(t) + sess, err := agent.New( + agent.WithModel("openrouter/free"), + agent.WithTools("read", "grep"), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + defer sess.Close() + + if got := strings.Join(sess.ToolNames(), ","); got != "read,grep" { + t.Errorf("ToolNames() = %q, want %q", got, "read,grep") + } +} + +// TestDenyWinsOverAllow is the fail-closed guarantee at the SDK layer: a tool on +// both lists is removed. +func TestDenyWinsOverAllow(t *testing.T) { + hermetic(t) + sess, err := agent.New( + agent.WithModel("openrouter/free"), + agent.WithTools("read", "bash"), + agent.WithDisallowedTools("bash"), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + defer sess.Close() + + if contains(sess.ToolNames(), "bash") { + t.Errorf("bash was on both lists and must be removed: %q", sess.ToolNames()) + } + if !contains(sess.ToolNames(), "read") { + t.Errorf("read was allowed and not denied, so must survive: %q", sess.ToolNames()) + } +} + +// TestWithoutTools yields an empty set — a pure text completion. +func TestWithoutTools(t *testing.T) { + hermetic(t) + sess, err := agent.New( + agent.WithModel("openrouter/free"), + agent.WithoutTools(), + ) + if err != nil { + t.Fatalf("New: %v", err) + } + defer sess.Close() + + if len(sess.ToolNames()) != 0 { + t.Errorf("WithoutTools must leave no tools, got %q", sess.ToolNames()) + } +} + +// TestUnknownToolIsError confirms a misspelled tool name fails construction +// rather than silently dropping the boundary. +func TestUnknownToolIsError(t *testing.T) { + hermetic(t) + _, err := agent.New( + agent.WithModel("openrouter/free"), + agent.WithTools("raed"), + ) + if err == nil { + t.Fatal("New = nil error, want a failure for the misspelled tool name") + } +} + +// TestInvalidThinkingLevelIsError confirms the level is validated up front. +func TestInvalidThinkingLevelIsError(t *testing.T) { + hermetic(t) + _, err := agent.New( + agent.WithModel("openrouter/free"), + agent.WithThinkingLevel("supersonic"), + ) + if err == nil { + t.Fatal("New = nil error, want a failure for the invalid thinking level") + } +} + +// TestValidThinkingLevels accepts every documented level. +func TestValidThinkingLevels(t *testing.T) { + hermetic(t) + for _, level := range []string{"off", "minimal", "low", "medium", "high", "xhigh", "max"} { + sess, err := agent.New( + agent.WithModel("openrouter/free"), + agent.WithThinkingLevel(level), + ) + if err != nil { + t.Errorf("WithThinkingLevel(%q): %v", level, err) + continue + } + sess.Close() + } +} + +// TestCloseHermetic confirms Close is a no-op (nil) when the session holds no +// plugin manager or memory store, and is safe to call. +func TestCloseHermetic(t *testing.T) { + hermetic(t) + sess, err := agent.New(agent.WithModel("openrouter/free"), agent.WithoutTools()) + if err != nil { + t.Fatalf("New: %v", err) + } + if err := sess.Close(); err != nil { + t.Errorf("Close() = %v, want nil", err) + } +} diff --git a/pigo/agent/doc.go b/pigo/agent/doc.go new file mode 100644 index 0000000..8279d97 --- /dev/null +++ b/pigo/agent/doc.go @@ -0,0 +1,63 @@ +// Package agent is the public, embeddable SDK for driving a pigo agent from +// your own Go program. It wraps pigo's internal run-assembly, provider, and +// agent-loop packages behind a small surface whose every exported type is a Go +// primitive (string, []string, bool, func) — so importing this package never +// pulls an internal type into your code, and pigo can evolve its internals +// without breaking you. +// +// # Quick start +// +// sess, err := agent.New( +// agent.WithModel("claude-opus-4-8"), +// agent.WithAPIKey(os.Getenv("ANTHROPIC_API_KEY")), +// ) +// if err != nil { +// log.Fatal(err) +// } +// defer sess.Close() +// +// reply, err := sess.Prompt(context.Background(), "Say hello in one word.") +// fmt.Println(reply) +// +// # Model, provider, credentials +// +// The model id selects the provider the same way the pigo CLI does: +// "claude-opus-4-8" resolves to Anthropic, "openrouter/free" to OpenRouter, +// and so on. Point at any OpenAI- or Anthropic-compatible endpoint with +// [WithBaseURL] + [WithProtocol], or a named provider from your config with +// [WithProvider]. The API key comes from [WithAPIKey] or, if unset, the +// provider's usual environment variable (e.g. ANTHROPIC_API_KEY). Keys are +// never logged. +// +// # Tools run automatically — read this +// +// By default a session is created with pigo's full built-in tool set (read, +// write, edit, bash, find, grep, and more) and those tools are executed WITHOUT +// any per-call confirmation prompt — equivalent to running the CLI with +// --approve. An agent can therefore read, modify, and delete files under its +// working directory and run shell commands on the host. This is the right +// default for an automated SDK, but it means you should only send prompts you +// trust, and run in a directory (and, ideally, a sandbox) you are willing to let +// the agent modify. To constrain or remove that capability use [WithTools] (an +// allowlist), [WithDisallowedTools] (a denylist, which always wins), or +// [WithoutTools] (a pure text completion with no tools at all). +// +// # Conversation state +// +// A [Session] keeps the running conversation: each [Session.Prompt] or +// [Session.Stream] call appends to the same history, so follow-up prompts see +// what came before. Call [Session.Reset] to start a fresh conversation on the +// same session, or [Session.Close] when you are done. A Session is NOT safe for +// concurrent use — drive it from one goroutine, or create one Session per +// goroutine. +// +// # Defaults +// +// - Tools: on (full built-in set, auto-executed; see the safety note above). +// - Skills: off — enable discovery of on-disk skills with [WithSkills]. +// - Memory: off — enable the persistent memory store with [WithMemory]. +// - Thinking: "medium" — override with [WithThinkingLevel]. +// +// Skills and memory are off by default so an embedded session is hermetic: it +// does not read or write the machine's shared pigo state unless you ask it to. +package agent diff --git a/pigo/agent/options.go b/pigo/agent/options.go new file mode 100644 index 0000000..a49ea58 --- /dev/null +++ b/pigo/agent/options.go @@ -0,0 +1,123 @@ +package agent + +// config is the resolved, unexported construction state for a Session. It is +// populated only through Option values, so callers never name or mutate it +// directly — the exported surface stays limited to With* constructors and the +// Session methods. +type config struct { + model string + baseURL string + protocol string + provider string + apiKey string + systemPrompt string + appendSystemPrompt []string + thinking string + noTools bool + allowedTools []string + disallowedTools []string + skills bool + memory bool +} + +// Option configures a Session at construction time. Options are applied in the +// order passed to New, so a later option overrides an earlier one that sets the +// same field. Because config is unexported, the only way to produce an Option is +// through the With* constructors below — which keeps the public surface free of +// internal types. +type Option func(*config) + +// WithModel sets the model id, which also selects the provider the way the pigo +// CLI does (e.g. "claude-opus-4-8" → Anthropic, "openrouter/free" → OpenRouter). +// The default is "openrouter/free". +func WithModel(model string) Option { + return func(c *config) { c.model = model } +} + +// WithBaseURL points the session at a custom endpoint. Pair it with +// [WithProtocol] to say whether that endpoint speaks the OpenAI or Anthropic +// wire format. +func WithBaseURL(baseURL string) Option { + return func(c *config) { c.baseURL = baseURL } +} + +// WithProtocol selects the wire protocol for a custom endpoint: "openai" or +// "anthropic". It is only consulted when [WithBaseURL] is set. +func WithProtocol(protocol string) Option { + return func(c *config) { c.protocol = protocol } +} + +// WithProvider selects a named provider from your pigo configuration instead of +// inferring one from the model id. +func WithProvider(name string) Option { + return func(c *config) { c.provider = name } +} + +// WithAPIKey sets the API key for the resolved provider, overriding the +// provider's environment variable. When unset, the provider's usual environment +// variable is used (e.g. ANTHROPIC_API_KEY, OPENROUTER_API_KEY). +func WithAPIKey(key string) Option { + return func(c *config) { c.apiKey = key } +} + +// WithSystemPrompt replaces pigo's built-in base instruction with prompt. Use +// this for full control over the agent's persona and rules; use +// [WithAppendSystemPrompt] instead to keep the built-in instruction and add to +// it. +func WithSystemPrompt(prompt string) Option { + return func(c *config) { c.systemPrompt = prompt } +} + +// WithAppendSystemPrompt appends one or more blocks to the system prompt, +// leaving pigo's built-in instruction in place. Repeated calls accumulate. +func WithAppendSystemPrompt(blocks ...string) Option { + return func(c *config) { + c.appendSystemPrompt = append(c.appendSystemPrompt, blocks...) + } +} + +// WithThinkingLevel sets the reasoning-effort level. Valid values are "off", +// "minimal", "low", "medium", "high", "xhigh", and "max". The default is +// "medium". An invalid value makes New return an error. +func WithThinkingLevel(level string) Option { + return func(c *config) { c.thinking = level } +} + +// WithTools restricts the session to the named built-in tools (an allowlist, +// e.g. WithTools("read", "grep")). Names are matched case-insensitively, so +// "Read" and "read" are equivalent. A name that matches no tool makes New +// return an error rather than silently ignoring it. Combine with +// [WithDisallowedTools]; deny always wins over allow. +func WithTools(names ...string) Option { + return func(c *config) { c.allowedTools = append(c.allowedTools, names...) } +} + +// WithDisallowedTools removes the named built-in tools (a denylist, e.g. +// WithDisallowedTools("bash")). Deny always wins: a tool named here is removed +// even if it also appears in [WithTools]. As with WithTools, an unknown name +// makes New return an error. +func WithDisallowedTools(names ...string) Option { + return func(c *config) { c.disallowedTools = append(c.disallowedTools, names...) } +} + +// WithoutTools removes every tool, producing a pure text-completion session that +// cannot touch the filesystem or run commands. It overrides [WithTools] and +// [WithDisallowedTools], which become inert once the tool set is empty. +func WithoutTools() Option { + return func(c *config) { c.noTools = true } +} + +// WithSkills enables discovery of on-disk skills, which are advertised to the +// model and loadable during a run. Skills are off by default so an embedded +// session stays independent of the machine's shared skills directory. +func WithSkills() Option { + return func(c *config) { c.skills = true } +} + +// WithMemory enables pigo's persistent memory store, letting the agent recall +// context saved by earlier runs and record new memories. Memory is off by +// default so an embedded session does not read or write shared state unless +// asked. +func WithMemory() Option { + return func(c *config) { c.memory = true } +} diff --git a/pigo/cmd/pigo/main.go b/pigo/cmd/pigo/main.go new file mode 100644 index 0000000..834bf74 --- /dev/null +++ b/pigo/cmd/pigo/main.go @@ -0,0 +1,548 @@ +// Command pigo is the CLI entry point for the pigo agent. It parses flags, +// overlays config.toml, and dispatches to one of the run modes — interactive +// REPL, headless print, session listing, or the internal sub-agent RPC server: +// +// pigo # interactive REPL (on a TTY) +// pigo -p "read README and summarize" # print mode: final text +// pigo -p "..." --output-format stream-json # line-delimited JSON events +// pigo install | list | uninstall | update # package management +// +// The provider is resolved from --model against the built-in OpenAI-compatible +// gateways (OpenRouter by default, Ollama for local models), with the API key +// taken from the environment. The process exit code reflects success (0) or +// failure (1), so the command composes cleanly in pipelines. All run-assembly, +// REPL, headless, and config logic lives under internal/cli/*; this file keeps +// only flag parsing (cliOptions), config overlay (applyFileConfig), and the +// dispatch seam that wires those subpackages together. +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "strings" + + flag "github.com/spf13/pflag" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/config" + "github.com/smallnest/pigo/internal/cli/headless" + "github.com/smallnest/pigo/internal/cli/pkgcmd" + "github.com/smallnest/pigo/internal/cli/repl" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/tui" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/dream" + "github.com/smallnest/pigo/internal/selfupdate" +) + +// Build metadata, injected at release time via -ldflags by goreleaser +// (see .goreleaser.yaml). They keep their default values for `go build`/ +// `go run` from source, so `pigo --version` still works without a release build. +var ( + version = "dev" + commit = "none" + date = "unknown" +) + +// cliOptions is the parsed command line, produced by main() and consumed by +// dispatch. Separating parse from dispatch makes the dispatch logic testable +// without touching the global flag set. +type cliOptions struct { + prompt string + model string + baseURL string + apiKey string + protocol string + // provider, when non-empty, selects a built-in provider by name from the + // registry (mirrors pi's provider selection): provider.ResolveProvider then builds the + // matching wire driver using the provider's default base URL, protocol, and + // API-key env var, ignoring the model-id heuristics. + provider string + outputFmt string + noTools bool + listSessions bool + resumeID string + continueLast bool + // approve grants the launch directory session-level trust up front (mirrors pi's + // --approve/-a): the first-launch trust prompt is skipped and side-effect + // tools (bash/write/edit) run without per-call confirmation for this run. + approve bool + // noSkills disables skill discovery (mirrors pi's --no-skills): skills under + // ~/.agents/skills are not loaded as /skill-name commands. + noSkills bool + // systemPrompt, when non-empty, replaces the default coding-assistant base + // instruction (mirrors pi's --system-prompt). The environment block and + // AGENTS.md injection still apply on top of it. + systemPrompt string + // appendSystemPrompt holds --append-system-prompt values (mirrors pi, repeatable): + // each is a path to a file whose contents are appended, or literal text when + // it is not an existing file. Appended after the base prompt and AGENTS.md. + appendSystemPrompt []string + // configPrompts holds prompt-template paths from the config.toml `prompts` + // array (settings tier); each is a file or directory loaded non-recursively. + // Populated by applyFileConfig; empty when the config omits `prompts`. + configPrompts []string + // promptTemplates holds --prompt-template paths (CLI tier, repeatable); each + // is a file or directory loaded non-recursively. + promptTemplates []string + // noPromptTemplates disables all prompt-template discovery (global, project, + // settings, CLI); built-in slash commands are unaffected. Independent of + // --no-skills. + noPromptTemplates bool + // subagentRPC selects the process-isolated sub-agent server mode (US-019, + // #135): pigo reads JSON-RPC sub-agent run requests from stdin and writes + // results to stdout. Internal, used by SubAgentTool's process mode. + subagentRPC bool + // dream, when set, runs the process-isolated memory-consolidation pass and + // exits: pigo enumerates + consolidates the global/project memory scope, emits + // a single-line Report JSON on stdout, and exits 0/1. Internal, spawned by the + // dream scheduler (and usable headlessly by scripts). See internal/dream and + // SPEC §4.1/§4.2. + dream bool + // dreamDryRun pairs with --dream: analyze and report without writing files or + // updating dream state (the lock is still taken). SPEC §5.5 dry-run row. + dreamDryRun bool + // thinkingLevel, when non-empty, is the --thinking-level flag: the reasoning + // effort for requests (off|minimal|low|medium|high|xhigh|max). It is the highest- + // precedence layer in resolveThinkingLevel, overriding PIGO_THINKING_LEVEL, the + // config files, and the built-in default (medium). + thinkingLevel string + // showVersion prints build metadata (version/commit/date, injected at release + // time by goreleaser) and exits, without running the agent. + showVersion bool + // noTUI forces the line-based REPL instead of the full-screen TUI (US-001). + // When set — or when stdout is not a TTY — the no-prompt path falls back to + // repl.Run rather than launching tui.Run. + noTUI bool + // cwd, when non-empty, is the working directory pigo switches to before doing + // anything else (matches the Claude Agent SDK's cwd option / git -C). Every + // cwd-derived resolution — built-in tool file roots, project trust, hooks + // project dir, .pigo/ project config, git info, the status-bar path — reads + // os.Getwd(), so a single os.Chdir here makes all of them operate in the + // given directory. This is what makes pigo usable as an SDK backend that can + // be pointed at an arbitrary project root. + cwd string + // memory holds the resolved [memory]/[checkpoint]/[compaction] config tables + // (defaults applied, string forms parsed). These have no CLI flags — the + // config file is their only source — so applyFileConfig always populates this + // (defaults when the tables are absent) for downstream memory/checkpoint/ + // compaction wiring to consume. See config.MemorySettings. + memory config.MemorySettings + // dreamCfg is the resolved [dream] configuration (enabled / interval / + // recent-sessions), populated by applyFileConfig from the [dream] table with + // defaults applied. The interactive REPL consumes it to decide the startup + // background auto-consolidation (US-008). Like memory it has no CLI flags. + dreamCfg dream.Config + // allowedTools and disallowedTools are the --allowed-tools/--disallowed-tools + // values: the tool-level admission boundary for the run, filling the gap + // between "all tools" and --no-tools. Each is repeatable and each value may be + // comma-separated. Names match case-insensitively, and deny wins over allow + // when a name appears on both sides (fail-closed). The boundary is enforced at + // the tool-registration layer in run.SetupEnv, strictly before the + // BeforeToolCall confirmation gate, so --approve waives confirmation prompts + // but can never widen the boundary. + allowedTools []string + disallowedTools []string +} + +func main() { + // Package-management subcommands (pigo install|list|uninstall|update ...) are + // positional and distinct from the flag-driven agent modes, so peel them off + // before pflag parsing — the agent flags don't apply to them. + if len(os.Args) > 1 && pkgcmd.Subcommands[os.Args[1]] { + // `pigo update` routes by whether a positional package name follows it: + // none — or flags-only, e.g. `pigo update --check` — is binary self-update + // (#466: download the latest release and replace this binary); a package + // name stays package-update (handled by pkgcmd). This is the US-003 dispatch + // split, with updateIsSelfUpdate as the pure classifier so routing is + // unit-testable (TestUpdateIsSelfUpdate). + if os.Args[1] == "update" && updateIsSelfUpdate(os.Args[2:]) { + os.Exit(selfupdate.Run(context.Background(), version, os.Stdout, os.Stderr)) + } + os.Exit(pkgcmd.Run(os.Args[1], os.Args[2:], os.Stdout, os.Stderr)) + } + + var opts cliOptions + flag.StringVarP(&opts.prompt, "print", "p", "", "prompt to run in headless print mode") + flag.StringVarP(&opts.model, "model", "m", "openrouter/free", "model id to run against (a well-known model name like claude-opus-4-8 or deepseek-chat auto-selects its provider when --provider/--protocol/--base-url are unset)") + flag.StringVarP(&opts.baseURL, "base-url", "u", "", "override provider base URL (e.g. local Ollama)") + flag.StringVarP(&opts.apiKey, "api-key", "k", "", "API key for the resolved provider (overrides env/config; else _API_KEY)") + flag.StringVarP(&opts.protocol, "protocol", "P", "", "force wire protocol for a custom endpoint: openai | anthropic (default: inferred from model id)") + flag.StringVar(&opts.provider, "provider", "", "select a built-in provider by name (e.g. deepseek, minimax); uses its default base URL, protocol, and API-key env var (see --help provider list)") + flag.StringVarP(&opts.outputFmt, "output-format", "o", "text", "output format: text | stream-json") + flag.BoolVarP(&opts.noTools, "no-tools", "n", false, "disable the built-in file/shell tools") + flag.StringArrayVar(&opts.allowedTools, "allowed-tools", nil, "restrict the model to these tools (repeatable, comma-separated, case-insensitive); empty means no restriction and --disallowed-tools wins on conflict") + flag.StringArrayVar(&opts.disallowedTools, "disallowed-tools", nil, "remove these tools from the model's set (repeatable, comma-separated, case-insensitive); takes precedence over --allowed-tools") + flag.BoolVarP(&opts.listSessions, "list-sessions", "l", false, "list stored interactive sessions and exit") + flag.StringVarP(&opts.resumeID, "resume", "r", "", "resume the interactive session with this id") + flag.BoolVarP(&opts.continueLast, "continue", "c", false, "resume the most recent interactive session") + flag.BoolVarP(&opts.approve, "approve", "a", false, "trust the working directory for this run: skip the first-launch trust prompt and run side-effect tools without per-call confirmation") + flag.BoolVar(&opts.noSkills, "no-skills", false, "disable skill discovery (do not load skills under ~/.agents/skills as /skill-name commands)") + flag.BoolVar(&opts.noPromptTemplates, "no-prompt-templates", false, "disable prompt-template discovery (do not load ~/.pigo/{commands,prompts}, .pigo/prompts, config prompts, or --prompt-template); built-in slash commands are unaffected") + flag.StringVar(&opts.systemPrompt, "system-prompt", "", "system prompt to use instead of the default coding-assistant prompt (mirrors pi --system-prompt)") + flag.StringArrayVar(&opts.appendSystemPrompt, "append-system-prompt", nil, "append text or file contents to the system prompt; repeatable (mirrors pi --append-system-prompt)") + flag.StringArrayVar(&opts.promptTemplates, "prompt-template", nil, "load a prompt template from a file or directory (non-recursive); repeatable (mirrors pi --prompt-template)") + flag.StringVar(&opts.thinkingLevel, "thinking-level", "", "reasoning effort: off|minimal|low|medium|high|xhigh|max (overrides PIGO_THINKING_LEVEL and config; default medium)") + flag.BoolVar(&opts.subagentRPC, "subagent-rpc", false, "internal: run as a process-isolated sub-agent JSON-RPC server over stdio (US-019)") + flag.BoolVar(&opts.dream, "dream", false, "internal: run a memory-consolidation pass over the global/project memory scope, emit a Report JSON on stdout, and exit (SPEC §4.1)") + flag.BoolVar(&opts.dreamDryRun, "dream-dry-run", false, "internal: with --dream, analyze and report without writing files or updating dream state (SPEC §5.5)") + flag.BoolVar(&opts.noTUI, "no-tui", false, "use the line-based REPL instead of the full-screen TUI") + flag.StringVarP(&opts.cwd, "cwd", "C", "", "run as if pigo was started in this directory (matches the Claude Agent SDK's cwd; like git -C): tool file access, trust, hooks, and project config all resolve against it") + flag.BoolVarP(&opts.showVersion, "version", "v", false, "print version information and exit") + // Extend the default pflag usage with a "Supported providers" block so + // `--help` documents the values accepted by --provider (name → env var → + // default base URL → protocol). The list is derived from the provider + // registry, so it never drifts from the code. + flag.Usage = func() { + out := flag.CommandLine.Output() + fmt.Fprintf(out, "Usage of %s:\n", os.Args[0]) + flag.PrintDefaults() + cli.PrintProviderHelp(out) + } + flag.Parse() + + // --cwd switches the process working directory before anything cwd-derived is + // resolved (tool roots, trust, hooks, project config, git info). Doing it here + // — after parse, before config overlay and dispatch — means every downstream + // os.Getwd() sees the requested directory, so pigo behaves as if it had been + // launched there. A bad path is a usage error (exit 2) rather than a silent + // fall-through to the original directory. + if opts.cwd != "" { + if err := os.Chdir(opts.cwd); err != nil { + fmt.Fprintf(os.Stderr, "pigo: --cwd: %v\n", err) + os.Exit(2) + } + } + + // Overlay ~/.config/pigo/config.toml: file values replace built-in defaults, + // but any flag the user set on the command line still wins (CLI > file > + // default). A malformed file warns but does not abort — defaults apply. + if cfg, err := config.LoadFileConfig(config.FileConfigPath()); err != nil { + fmt.Fprintf(os.Stderr, "pigo: %v\n", err) + } else { + applyFileConfig(&opts, cfg, flag.CommandLine.Changed) + } + + // --version is a standalone action: print build metadata and exit. + if opts.showVersion { + fmt.Printf("pigo %s (commit %s, built %s)\n", version, commit, date) + os.Exit(0) + } + + // A prompt may also be supplied as positional args. + if opts.prompt == "" { + opts.prompt = strings.TrimSpace(strings.Join(flag.Args(), " ")) + } + + os.Exit(dispatch(context.Background(), opts, os.Stdout, os.Stderr)) +} + +// applyFileConfig overlays config.toml values onto opts, but only for flags the +// user did not set on the command line (changed reports whether a flag name was +// explicitly passed). This yields the precedence: CLI flag > config file > +// default. Zero-valued config fields never override. +func applyFileConfig(opts *cliOptions, cfg config.FileConfig, changed func(string) bool) { + if cfg.Model != "" && !changed("model") { + opts.model = cfg.Model + } + if cfg.BaseURL != "" && !changed("base-url") { + opts.baseURL = cfg.BaseURL + } + if cfg.APIKey != "" && !changed("api-key") { + opts.apiKey = cfg.APIKey + } + if cfg.Protocol != "" && !changed("protocol") { + opts.protocol = cfg.Protocol + } + if cfg.Provider != "" && !changed("provider") { + opts.provider = cfg.Provider + } + if cfg.ThinkingLevel != "" && !changed("thinking-level") { + opts.thinkingLevel = cfg.ThinkingLevel + } + if cfg.OutputFormat != "" && !changed("output-format") { + opts.outputFmt = cfg.OutputFormat + } + if cfg.NoTools && !changed("no-tools") { + opts.noTools = true + } + if cfg.NoSkills && !changed("no-skills") { + opts.noSkills = true + } + if cfg.Approve && !changed("approve") { + opts.approve = true + } + if cfg.SystemPrompt != "" && !changed("system-prompt") { + opts.systemPrompt = cfg.SystemPrompt + } + // The tool boundary follows the standard precedence (CLI > file > default) + // rather than the additive treatment prompts get below. Merging would be the + // wrong semantics for a security boundary: a user passing --allowed-tools to + // widen what the file's allowed_tools narrowed must actually get the wider + // set, not the intersection. Each flag overrides its own key independently: + // --allowed-tools does not clear a file-level disallowed_tools, and because + // deny wins on conflict a file deny survives a CLI allow — re-admitting a + // file-denied tool requires overriding --disallowed-tools on the CLI. + if len(cfg.AllowedTools) > 0 && !changed("allowed-tools") { + opts.allowedTools = cfg.AllowedTools + } + if len(cfg.DisallowedTools) > 0 && !changed("disallowed-tools") { + opts.disallowedTools = cfg.DisallowedTools + } + // prompts (settings tier) are additive with --prompt-template (CLI tier, + // wired in #339), so they are always passed through when present. + if len(cfg.Prompts) > 0 { + opts.configPrompts = cfg.Prompts + } + // The [memory]/[checkpoint]/[compaction] tables have no CLI flags, so they + // are resolved (with defaults) and overlaid unconditionally — an absent set + // of tables yields the default-safe MemorySettings. + opts.memory = cfg.ResolveMemorySettings() + // The [dream] table also has no CLI flags; normalize it (defaults applied when + // the table is absent) so the interactive startup trigger has a resolved + // Config. NewConfig treats a nil enabled as true, so dream is on by default. + opts.dreamCfg = dream.NewConfig(cfg.Dream.Enabled, cfg.Dream.IntervalDays, cfg.Dream.RecentSessions) +} + +// dispatch runs the resolved command and returns a process exit code, writing +// diagnostics to errOut. It is the run-assembly seam: every path (list, REPL, +// headless, subagent-rpc) is reached from here, so the CLI's behavior can be +// exercised without re-parsing flags. A returned code of 0 is success. +func dispatch(ctx context.Context, opts cliOptions, out, errOut io.Writer) int { + // --subagent-rpc is a fully separate mode: speak the sub-agent JSON-RPC + // protocol over stdio and exit. It is the subprocess end of process-isolated + // sub-agents and shares nothing with the interactive/headless paths. + if opts.subagentRPC { + return headless.RunSubAgentRPC(ctx, os.Stdin, out, errOut) + } + + // --dream is the subprocess consolidation mode (SPEC §4.1/§4.2): run one + // memory-consolidation pass to completion, emit a single-line Report JSON on + // stdout (progress/logs go to stderr), and exit 0 on success / 1 on failure. + // It runs before any interactive/headless session assembly and honors -C/--cwd + // for the project scope (applied above via os.Chdir). It shares nothing with + // the REPL/headless paths. + if opts.dream { + return runDream(ctx, opts, out, errOut) + } + + // --list-sessions is a standalone action: print and exit. + if opts.listSessions { + if err := headless.PrintSessions(out); err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + return 0 + } + + // --continue resolves to the most recently updated session id. + resumeID := opts.resumeID + if opts.continueLast && resumeID == "" { + id, err := headless.MostRecentSessionID() + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + if id == "" { + fmt.Fprintln(errOut, "pigo: no sessions to continue") + return 1 + } + resumeID = id + } + + // No prompt + an interactive terminal → start the interactive UI. By default + // this is the full-screen TUI (US-001); --no-tui (or a non-terminal stdout) + // forces the line-based REPL (US-003). A --resume id also enters the + // interactive UI to continue an existing session. No prompt with a + // non-terminal stdout (pipe/CI) and no resume is an error, since there is + // nothing to run and nothing to interact with. + if opts.prompt == "" { + isTTY := ui.StdoutIsTerminal() + if resumeID == "" && !isTTY { + fmt.Fprintln(errOut, "pigo: no prompt (use -p \"...\" or positional args)") + return 2 + } + env, err := run.SetupEnv(opts.model, opts.baseURL, opts.protocol, opts.provider, opts.apiKey, opts.noTools, opts.noSkills, opts.systemPrompt, opts.appendSystemPrompt, opts.memory.Memory.Enabled, run.NewToolPolicy(opts.allowedTools, opts.disallowedTools)) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return setupExitCode(err) + } + if env.Plugins != nil { + defer env.Plugins.Close() + } + if env.Memory != nil { + defer env.Memory.Close() + } + thinking, err := run.ResolveThinkingLevel(opts.thinkingLevel) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 2 + } + if shouldUseTUI(opts, isTTY) { + // Refresh the cached latest-release check off the hot path so the banner + // can show an upgrade hint on this or the next launch without blocking + // startup (US-004). No-ops for dev builds or a fresh cache. + selfupdate.StartBackgroundCheck(version) + if err := tui.Run(tui.Options{ + Model: opts.model, + ProviderName: env.ProviderName, + Provider: env.Provider, + BaseURL: opts.baseURL, + APIKey: opts.apiKey, + Protocol: opts.protocol, + Version: version, + ThinkingLevel: thinking, + Tools: env.Tools, + SysPrompt: env.SysPrompt, + ResumeID: resumeID, + Approve: opts.approve, + Skills: env.Skills, + Plugins: env.Plugins, + ConfigPrompts: opts.configPrompts, + CliPrompts: opts.promptTemplates, + NoPromptTemplates: opts.noPromptTemplates, + }); err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + return 0 + } + if err := repl.Run(repl.Options{ + Model: opts.model, + ProviderName: env.ProviderName, + Provider: env.Provider, + BaseURL: opts.baseURL, + APIKey: opts.apiKey, + Protocol: opts.protocol, + ThinkingLevel: thinking, + Tools: env.Tools, + SysPrompt: env.SysPrompt, + ResumeID: resumeID, + Approve: opts.approve, + Skills: env.Skills, + Plugins: env.Plugins, + ConfigPrompts: opts.configPrompts, + CliPrompts: opts.promptTemplates, + NoPromptTemplates: opts.noPromptTemplates, + Dream: opts.dreamCfg, + }); err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + return 0 + } + + mode, err := headless.ParseOutputMode(opts.outputFmt) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 2 + } + + env, err := run.SetupEnv(opts.model, opts.baseURL, opts.protocol, opts.provider, opts.apiKey, opts.noTools, opts.noSkills, opts.systemPrompt, opts.appendSystemPrompt, opts.memory.Memory.Enabled, run.NewToolPolicy(opts.allowedTools, opts.disallowedTools)) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return setupExitCode(err) + } + if env.Plugins != nil { + defer env.Plugins.Close() + } + if env.Memory != nil { + defer env.Memory.Close() + } + return headless.Run(ctx, headless.RunParams{ + Mode: mode, + Env: env, + Prompt: opts.prompt, + Model: opts.model, + APIKey: opts.apiKey, + ThinkingLevel: opts.thinkingLevel, + ResumeID: resumeID, + }, out, errOut) +} + +// setupExitCode maps a run.SetupEnv failure to a process exit code. A bad tool +// policy is a usage error (2), matching --cwd and --output-format; everything +// else — provider resolution, prompt assembly — is a runtime failure (1). +func setupExitCode(err error) int { + var policyErr *run.ToolPolicyError + if errors.As(err, &policyErr) { + return 2 + } + return 1 +} + +// runDream executes the subprocess memory-consolidation pass (SPEC §4.1/§4.2). +// It runs dream.Runner to completion, marshals the resulting Report as a single +// line of JSON on stdout (the parent/scheduler parses this), and returns the +// process exit code: 0 on success (including a "skipped" run when another dream +// holds the lock) or 1 on failure. Progress and diagnostics go to errOut. The +// project scope comes from the working directory, which -C/--cwd already applied +// via os.Chdir before dispatch, so an empty ProjectDir here resolves to cwd. +func runDream(ctx context.Context, opts cliOptions, out, errOut io.Writer) int { + projectDir, err := os.Getwd() + if err != nil { + fmt.Fprintf(errOut, "pigo: dream: %v\n", err) + return 1 + } + // The dream pass reuses the main-session model (SPEC Q3): resolve the same + // model/provider/api-key tuple cmd/pigo already overlaid from flags+config, + // and inject a real LLM-backed Consolidator so `pigo --dream` performs the + // semantic merge/prune step (not just the deterministic dedup/path-clean). + thinking, err := run.ResolveThinkingLevel(opts.thinkingLevel) + if err != nil { + fmt.Fprintf(errOut, "pigo: dream: %v\n", err) + return 1 + } + cons, err := dream.NewLLMConsolidator(opts.model, opts.baseURL, opts.protocol, opts.provider, opts.apiKey, thinking) + if err != nil { + fmt.Fprintf(errOut, "pigo: dream: %v\n", err) + return 1 + } + r := &dream.Runner{Consolidator: cons} + report, err := r.Run(ctx, dream.RunOptions{ + DryRun: opts.dreamDryRun, + ProjectDir: projectDir, + }) + if err != nil { + fmt.Fprintf(errOut, "pigo: dream: %v\n", err) + return 1 + } + // Single-line JSON on stdout is the stdout contract (SPEC §4.2). Encoder + // writes a trailing newline, keeping the report one line. + if err := json.NewEncoder(out).Encode(report); err != nil { + fmt.Fprintf(errOut, "pigo: dream: encode report: %v\n", err) + return 1 + } + return 0 +} + +// shouldUseTUI is the pure entry-gating predicate for the no-prompt path +// (US-001, SPEC 4.2/5.2): the full-screen TUI is used only when stdout is a TTY +// and --no-tui was not set. --no-tui or a non-terminal stdout always forces the +// line-based REPL. Keeping the decision in a side-effect-free function lets the +// gating be unit-tested without a real terminal or spawning Bubble Tea (see +// TestDispatchTUIGating); dispatch handles the non-TTY/no-resume usage error +// before calling this, so it only decides TUI-vs-REPL for the interactive case. +func shouldUseTUI(opts cliOptions, isTTY bool) bool { + return isTTY && !opts.noTUI +} + +// updateIsSelfUpdate classifies the arguments that follow `pigo update` (US-003) +// to route between binary self-update and pkgmgr package-update. It returns true +// — self-update — when no positional package name is present: any argument that +// does not begin with '-' is treated as a package name and routes to +// package-update, while flags-only invocations (e.g. `pigo update --check`) stay +// on the self-update path. Keeping the decision side-effect-free lets the routing +// be unit-tested without spawning either update path (see TestUpdateIsSelfUpdate). +func updateIsSelfUpdate(rest []string) bool { + for _, a := range rest { + if !strings.HasPrefix(a, "-") { + return false + } + } + return true +} diff --git a/pigo/cmd/pigo/main_test.go b/pigo/cmd/pigo/main_test.go new file mode 100644 index 0000000..eada7e0 --- /dev/null +++ b/pigo/cmd/pigo/main_test.go @@ -0,0 +1,350 @@ +package main + +// Tests for the thin CLI entry point: the dispatch seam (options+writers → +// exit code), the config.toml overlay (applyFileConfig precedence), and the +// settings-tier prompts pass-through. These exercise the branching without +// spawning a provider or re-parsing the global flag set. + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/cli/config" + "github.com/smallnest/pigo/internal/cli/run" +) + +// --- dispatch seam --- + +// TestDispatchListSessionsEmpty verifies --list-sessions is a standalone action +// that succeeds (exit 0) and prints the empty-store message, using an isolated +// PIGO_HOME so it never touches the real session store. +func TestDispatchListSessionsEmpty(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + var out, errOut bytes.Buffer + code := dispatch(context.Background(), cliOptions{listSessions: true}, &out, &errOut) + if code != 0 { + t.Fatalf("exit code = %d, want 0 (errOut=%q)", code, errOut.String()) + } + if !strings.Contains(out.String(), "no sessions") { + t.Errorf("out = %q, want the empty-store message", out.String()) + } +} + +// TestDispatchContinueNoSessions verifies --continue with an empty store is an +// error (exit 1) that says there is nothing to continue, rather than starting a +// blank REPL. +func TestDispatchContinueNoSessions(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + // Ensure a non-terminal path is not taken before the continue guard: continue + // resolves the id first and errors when the store is empty. + var out, errOut bytes.Buffer + code := dispatch(context.Background(), cliOptions{continueLast: true}, &out, &errOut) + if code != 1 { + t.Fatalf("exit code = %d, want 1 (errOut=%q)", code, errOut.String()) + } + if !strings.Contains(errOut.String(), "no sessions to continue") { + t.Errorf("errOut = %q, want the no-sessions-to-continue message", errOut.String()) + } +} + +// TestDispatchNoPromptNonTerminal verifies the CI/pipe guard: no prompt, no +// resume, and a non-terminal stdout is a usage error (exit 2) with a diagnostic +// on errOut — reachable now that dispatch takes its writers as parameters. +func TestDispatchNoPromptNonTerminal(t *testing.T) { + var out, errOut bytes.Buffer + code := dispatch(context.Background(), cliOptions{}, &out, &errOut) + if code != 2 { + t.Fatalf("exit code = %d, want 2", code) + } + if !strings.Contains(errOut.String(), "no prompt") { + t.Errorf("errOut = %q, want it to mention the missing prompt", errOut.String()) + } +} + +// TestDispatchBadOutputFormat verifies an unknown --output-format is rejected +// (exit 2) before any provider work, naming the offending value. +func TestDispatchBadOutputFormat(t *testing.T) { + var out, errOut bytes.Buffer + code := dispatch(context.Background(), cliOptions{prompt: "hi", outputFmt: "yaml"}, &out, &errOut) + if code != 2 { + t.Fatalf("exit code = %d, want 2", code) + } + if !strings.Contains(errOut.String(), "yaml") { + t.Errorf("errOut = %q, want it to name the bad format", errOut.String()) + } +} + +// TestDispatchTUIGating verifies the pure entry-gating predicate shouldUseTUI +// (US-001, SPEC 4.2/5.2) that dispatch uses to choose the full-screen TUI vs the +// line-based REPL on the no-prompt path. The decision is tested directly so it +// needs no real TTY and never spawns Bubble Tea: TUI only when stdout is a TTY +// and --no-tui is unset; --no-tui or a non-terminal stdout always forces REPL. +func TestDispatchTUIGating(t *testing.T) { + tests := []struct { + name string + opts cliOptions + isTTY bool + want bool + }{ + {name: "TTY and no flag uses TUI", opts: cliOptions{}, isTTY: true, want: true}, + {name: "--no-tui forces REPL on a TTY", opts: cliOptions{noTUI: true}, isTTY: true, want: false}, + {name: "non-TTY never uses TUI", opts: cliOptions{}, isTTY: false, want: false}, + {name: "non-TTY with --no-tui stays REPL", opts: cliOptions{noTUI: true}, isTTY: false, want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := shouldUseTUI(tt.opts, tt.isTTY); got != tt.want { + t.Errorf("shouldUseTUI(%+v, %v) = %v, want %v", tt.opts, tt.isTTY, got, tt.want) + } + }) + } +} + +// TestCwdChdirRootsEnv verifies the guarantee --cwd relies on: after the +// process working directory is switched (what the --cwd flag does via os.Chdir), +// run.SetupEnv roots the run — and thus the built-in file tools — at that +// directory. This is the contract that lets pigo be pointed at an arbitrary +// project root as an SDK backend. It exercises the downstream effect rather than +// re-parsing flags, since the chdir itself lives in main(). +func TestCwdChdirRootsEnv(t *testing.T) { + orig, err := os.Getwd() + if err != nil { + t.Fatalf("Getwd: %v", err) + } + t.Cleanup(func() { _ = os.Chdir(orig) }) + + dir := t.TempDir() + if err := os.Chdir(dir); err != nil { + t.Fatalf("Chdir: %v", err) + } + // macOS temp dirs are symlinks (/tmp → /private/tmp); os.Getwd resolves them, + // so compare against the resolved form. + want, err := filepath.EvalSymlinks(dir) + if err != nil { + t.Fatalf("EvalSymlinks: %v", err) + } + + env, err := run.SetupEnv("openrouter/free", "", "", "", "", true /*noTools*/, true /*noSkills*/, "", nil, false /*memEnabled*/, run.ToolPolicy{}) + if err != nil { + t.Fatalf("SetupEnv: %v", err) + } + if env.Cwd != want { + t.Errorf("env.Cwd = %q, want %q (the chdir'd directory)", env.Cwd, want) + } +} + +// TestUpdateIsSelfUpdate verifies the US-003 `pigo update` routing classifier: +// no positional package name (including flags-only invocations like +// `pigo update --check`) routes to binary self-update (true); any positional +// package name routes to pkgmgr package-update (false). Tested directly so the +// dispatch split needs no argv parsing or spawning either update path. +func TestUpdateIsSelfUpdate(t *testing.T) { + tests := []struct { + name string + rest []string + want bool + }{ + {name: "no args is self-update", rest: nil, want: true}, + {name: "empty slice is self-update", rest: []string{}, want: true}, + {name: "flags-only is self-update", rest: []string{"--check"}, want: true}, + {name: "multiple flags is self-update", rest: []string{"--check", "-v"}, want: true}, + {name: "single package name is package-update", rest: []string{"pi-mcp-adapter"}, want: false}, + {name: "multiple package names is package-update", rest: []string{"a", "b"}, want: false}, + {name: "flag then package name is package-update", rest: []string{"--check", "pkg"}, want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := updateIsSelfUpdate(tt.rest); got != tt.want { + t.Errorf("updateIsSelfUpdate(%v) = %v, want %v", tt.rest, got, tt.want) + } + }) + } +} + +// --- config.toml overlay --- + +// changedSet turns a set of flag names into a lookup func for applyFileConfig. +func changedSet(names ...string) func(string) bool { + set := make(map[string]bool, len(names)) + for _, n := range names { + set[n] = true + } + return func(name string) bool { return set[name] } +} + +func TestApplyFileConfig_FillsUnsetFlags(t *testing.T) { + opts := cliOptions{model: "openrouter/free", outputFmt: "text"} + cfg := config.FileConfig{ + Model: "claude-opus-4-8", + BaseURL: "https://example.com", + APIKey: "sk-test", + Protocol: "anthropic", + Provider: "deepseek", + ThinkingLevel: "high", + OutputFormat: "stream-json", + NoTools: true, + NoSkills: true, + Approve: true, + SystemPrompt: "be terse", + } + applyFileConfig(&opts, cfg, changedSet()) + + if opts.model != "claude-opus-4-8" { + t.Errorf("model = %q, want claude-opus-4-8", opts.model) + } + if opts.baseURL != "https://example.com" { + t.Errorf("baseURL = %q", opts.baseURL) + } + if opts.apiKey != "sk-test" { + t.Errorf("apiKey = %q", opts.apiKey) + } + if opts.protocol != "anthropic" { + t.Errorf("protocol = %q", opts.protocol) + } + if opts.provider != "deepseek" { + t.Errorf("provider = %q", opts.provider) + } + if opts.thinkingLevel != "high" { + t.Errorf("thinkingLevel = %q", opts.thinkingLevel) + } + if opts.outputFmt != "stream-json" { + t.Errorf("outputFmt = %q", opts.outputFmt) + } + if !opts.noTools || !opts.noSkills || !opts.approve { + t.Errorf("bool flags not applied: %+v", opts) + } + if opts.systemPrompt != "be terse" { + t.Errorf("systemPrompt = %q", opts.systemPrompt) + } +} + +func TestApplyFileConfig_CLIWins(t *testing.T) { + opts := cliOptions{model: "cli-model", outputFmt: "text"} + cfg := config.FileConfig{Model: "config-model", OutputFormat: "stream-json"} + // --model was set on the command line; --output-format was not. + applyFileConfig(&opts, cfg, changedSet("model")) + + if opts.model != "cli-model" { + t.Errorf("CLI model should win, got %q", opts.model) + } + if opts.outputFmt != "stream-json" { + t.Errorf("unset output-format should take config value, got %q", opts.outputFmt) + } +} + +func TestApplyFileConfig_EmptyConfigNoChange(t *testing.T) { + opts := cliOptions{model: "openrouter/free", outputFmt: "text"} + applyFileConfig(&opts, config.FileConfig{}, changedSet()) + if opts.model != "openrouter/free" || opts.outputFmt != "text" { + t.Fatalf("empty config should not change opts, got %+v", opts) + } + if opts.baseURL != "" || opts.provider != "" || opts.noTools { + t.Fatalf("empty config should leave unset fields empty, got %+v", opts) + } +} + +func TestApplyFileConfigPrompts(t *testing.T) { + var opts cliOptions + cfg := config.FileConfig{Prompts: []string{"./my-prompts", "/abs/x.md"}} + applyFileConfig(&opts, cfg, func(string) bool { return false }) + if len(opts.configPrompts) != 2 || opts.configPrompts[0] != "./my-prompts" || opts.configPrompts[1] != "/abs/x.md" { + t.Errorf("opts.configPrompts = %v, want [./my-prompts /abs/x.md]", opts.configPrompts) + } +} + +// The tool boundary follows CLI > file > default like the other scalar flags, +// rather than the additive treatment `prompts` gets: merging would prevent a CLI +// flag from widening a boundary the config file narrowed. +func TestApplyFileConfigToolPolicy(t *testing.T) { + t.Run("fills unset flags", func(t *testing.T) { + var opts cliOptions + cfg := config.FileConfig{ + AllowedTools: []string{"read", "grep"}, + DisallowedTools: []string{"bash"}, + } + applyFileConfig(&opts, cfg, changedSet()) + if len(opts.allowedTools) != 2 || opts.allowedTools[0] != "read" { + t.Errorf("allowedTools = %v, want [read grep]", opts.allowedTools) + } + if len(opts.disallowedTools) != 1 || opts.disallowedTools[0] != "bash" { + t.Errorf("disallowedTools = %v, want [bash]", opts.disallowedTools) + } + }) + + t.Run("CLI replaces file value wholesale", func(t *testing.T) { + opts := cliOptions{allowedTools: []string{"bash"}} + cfg := config.FileConfig{AllowedTools: []string{"read"}, DisallowedTools: []string{"write"}} + applyFileConfig(&opts, cfg, changedSet("allowed-tools")) + if len(opts.allowedTools) != 1 || opts.allowedTools[0] != "bash" { + t.Errorf("CLI --allowed-tools must win outright, got %v", opts.allowedTools) + } + if len(opts.disallowedTools) != 1 || opts.disallowedTools[0] != "write" { + t.Errorf("unset --disallowed-tools should take the config value, got %v", opts.disallowedTools) + } + }) + + t.Run("absent config leaves the boundary open", func(t *testing.T) { + var opts cliOptions + applyFileConfig(&opts, config.FileConfig{}, changedSet()) + if opts.allowedTools != nil || opts.disallowedTools != nil { + t.Errorf("empty config must not constrain tools, got %v / %v", opts.allowedTools, opts.disallowedTools) + } + }) +} + +// setupExitCode maps a bad tool policy to the usage exit code (2) and everything +// else to a runtime failure (1), so a typo is distinguishable from e.g. a +// provider-resolution error. +func TestSetupExitCode(t *testing.T) { + if got := setupExitCode(&run.ToolPolicyError{UnknownAllowed: []string{"raed"}}); got != 2 { + t.Errorf("setupExitCode(ToolPolicyError) = %d, want 2 (usage error)", got) + } + if got := setupExitCode(errors.New("provider boom")); got != 1 { + t.Errorf("setupExitCode(generic) = %d, want 1", got) + } + if got := setupExitCode(fmt.Errorf("wrapped: %w", &run.ToolPolicyError{})); got != 2 { + t.Errorf("setupExitCode must unwrap, got %d, want 2", got) + } +} + +// applyFileConfig always resolves the [memory]/[checkpoint]/[compaction] tables +// into opts.memory, applying defaults when they are absent. +func TestApplyFileConfig_MemoryDefaults(t *testing.T) { + var opts cliOptions + applyFileConfig(&opts, config.FileConfig{}, changedSet()) + if !opts.memory.Memory.Enabled || !opts.memory.Memory.ReconcileOnSearch { + t.Errorf("memory defaults not applied: %+v", opts.memory.Memory) + } + if opts.memory.Memory.SearchScoreFloor != 0.15 { + t.Errorf("search_score_floor default = %v, want 0.15", opts.memory.Memory.SearchScoreFloor) + } + if len(opts.memory.CheckpointThresholds) != 3 { + t.Errorf("checkpoint thresholds default = %v, want 3 entries", opts.memory.CheckpointThresholds) + } + if opts.memory.MaxContext.IsSet() { + t.Errorf("max_context should be unset by default") + } +} + +// A configured [memory]/[compaction] set overlays into opts.memory. +func TestApplyFileConfig_MemoryOverride(t *testing.T) { + var opts cliOptions + enabled := false + cfg := config.FileConfig{ + Memory: config.MemoryConfig{Enabled: &enabled}, + Compaction: config.CompactionConfig{MaxContext: "50%"}, + } + applyFileConfig(&opts, cfg, changedSet()) + if opts.memory.Memory.Enabled { + t.Errorf("memory.enabled=false should overlay, got true") + } + if got := opts.memory.MaxContext.Resolve(200000); got != 100000 { + t.Errorf("max_context 50%% of 200000 = %d, want 100000", got) + } +} diff --git a/pigo/config.toml.example b/pigo/config.toml.example new file mode 100644 index 0000000..5fd8aa3 --- /dev/null +++ b/pigo/config.toml.example @@ -0,0 +1,92 @@ +# Pigo global configuration (example template) +# Copy this file to ~/.config/pigo/config.toml and fill in your own values. +# https://github.com/smallnest/pigo + +# Provider base URL (overrides the provider's default endpoint). +# Leave unset to use the selected provider's built-in default. +# base_url = "https://api.anthropic.com/v1" + +# Default model to use. +model = "claude-opus-4-8" + +# API key for authentication. +# Prefer supplying this via the _API_KEY environment variable +# (e.g. ANTHROPIC_API_KEY) instead of committing it to a config file. +# api_key = "your-api-key-here" + +# Output format: text | stream-json +output_format = "text" + +# Trust the working directory for this run: skip the first-launch trust prompt. +approve = false + +# Disable the built-in file/shell tools. +no_tools = false + +# Tool-level admission control (the middle ground between "all tools" and +# no_tools). Names match case-insensitively, so "Read" hits the built-in "read". +# disallowed_tools wins when a name appears in both lists (fail-closed). +# +# A boundary declared here is a HARD boundary: filtering happens at the tool- +# registration layer, so `approve = true` waives per-call confirmation but can +# never let the model reach a tool outside the boundary. Task sub-agents inherit +# it too. An unknown tool name aborts startup with exit code 2 rather than being +# silently ignored. +# +# Each CLI flag REPLACES its own file value wholesale (it does not merge): +# --allowed-tools overrides allowed_tools, --disallowed-tools overrides +# disallowed_tools, independently. So a CLI --allowed-tools can widen what +# allowed_tools narrowed. Note the two lists stay independent and deny still +# wins: a file-level disallowed_tools is NOT lifted by a CLI --allowed-tools — +# to re-admit a tool the file denied, override --disallowed-tools on the CLI. +# Parameter-level forms such as Bash(git log:*) are not supported yet. +# allowed_tools = ["read", "grep"] +# disallowed_tools = ["bash", "bash_output", "kill_bash"] + +# Disable skill discovery (do not load skills under ~/.agents/skills as /skill-name commands). +no_skills = false + +# Force wire protocol for a custom endpoint: openai | anthropic +# protocol = "anthropic" + +# Run in headless print mode with the given prompt. +# print = "" + +# Resume the most recent interactive session. +# continue = false + +# List stored interactive sessions and exit. +# list_sessions = false + +# Resume the interactive session with this id. +# resume = "" + +# Internal: run as a process-isolated sub-agent JSON-RPC server over stdio (US-019). +# subagent_rpc = false + +# --- Persistent memory + infinite context (nested tables) --- + +# [memory] controls the persistent memory system. All keys are optional and +# default to the values shown; set memory.enabled = false to fully disable it. +# [memory] +# enabled = true # master switch for the memory system +# reconcile_on_search = true # lazily re-index memory files before each search +# search_score_floor = 0.15 # drop search hits below this relevance score (0..1) +# cc_index = false # also index the Claude Code memory directory (read-only) + +# [checkpoint] controls infinite-context checkpointing. +# [checkpoint] +# thresholds = ["40%", "60%", "80%"] # context-window fill levels that trigger compaction +# reserved = 4096 # optional: tokens to reserve (int) or a percentage ("10%") + +# [checkpoint.push_caps] caps per-section injection token budgets. +# [checkpoint.push_caps] +# memory = 800 +# recall = 1200 + +# [compaction] tunes auto-compaction. +# [compaction] +# max_context accepts a token count (300000), a K/M suffix ("300K", "1M"), or a +# percentage of the provider window ("50%"). It only lowers the trigger point and +# is always clamped by the provider limit. +# max_context = "300K" diff --git a/pigo/coop/AGENTS.md b/pigo/coop/AGENTS.md new file mode 100644 index 0000000..1dd0fc4 --- /dev/null +++ b/pigo/coop/AGENTS.md @@ -0,0 +1,29 @@ +# 单 agent 任务协议 + +你是独立的任务执行 agent,完成 /blackboard/task.md 中的任务。任务信息通过工作区与黑板目录交换。 + +## 目录结构 + +| 路径 | 用途 | +| --- | --- | +| /blackboard/task.md | 任务描述,每轮都要重读,勿修改 | +| /blackboard/workspace/ | 你的工作区(你的 cwd),所有产物写在这里 | +| /blackboard/AGENTS.md | 本协议(副本在你的工作区,勿修改) | +| /blackboard/DONE | 完成标记,存在即表示任务已交付完成,只用 blackboard 工具创建 | +| /blackboard/logs/ | 每轮运行日志(supervisor 维护) | +| /blackboard/sessions/ | 会话 ID(supervisor 维护,供 --resume) | +| /blackboard/result.json | 任务结果(supervisor 结束时生成,勿手动修改) | + +## 每轮流程 + +1. **读取**:读取 /blackboard/task.md,明确任务与提交规则(若为解题任务,通常含 flag 提交 API / token / unique_code)。 +2. **评估**:检查工作区已有产物,判断进度与缺口,不重复已完成工作。 +3. **执行**:推进任务——执行命令、读写文件、编写产物到 /blackboard/workspace/。 +4. **提交**:拿到 flag 的任务,立即按 task.md 中约定的规则提交,并把提交响应与得分写入产物。 +5. **判定**:全部工作完成、交付物完整后,用 blackboard 工具 action=done 创建 DONE(summary 为最终交付总结,含关键结果/flag/提交响应/产物清单)。DONE 只能创建一次;宁可多轮打磨,不要过早宣布完成;确认无法解出时同样用 done 标记并写明「未解出」与已尝试内容。 + +## 红线 + +- 不直接写 /blackboard/DONE——必须走 blackboard 工具 action=done,保证原子创建。 +- 不伪造实验、命令或提交结果;如实记录。 +- 不修改 /blackboard/task.md、/blackboard/AGENTS.md 与 /blackboard/result.json。 diff --git a/pigo/coop/Dockerfile b/pigo/coop/Dockerfile new file mode 100644 index 0000000..c6b3bd2 --- /dev/null +++ b/pigo/coop/Dockerfile @@ -0,0 +1,56 @@ +# pigo 协作镜像(单 agent 版) +# +# 构建(在仓库根目录执行,先本地交叉编译 pigo,再 docker build): +# $env:GOOS="linux"; $env:CGO_ENABLED="0" +# go build -trimpath -ldflags="-s -w" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo +# docker build -f coop/Dockerfile -t pigo-coop . +# 或直接用 Linux 的 Go 工具链时跳过交叉编译步骤。 +# +# 运行示例(模型后端由外部调用方通过环境变量传入): +# docker run --rm \ +# -e MODEL=deepseek-chat \ +# -e BASE_URL=https://api.deepseek.com \ +# -e API_KEY= \ +# -e TASK="请为 XX 编写设计文档并实现原型" \ +# -e ROUND_MAX=6 \ +# pigo-coop +# +# 协作方式:单个 pigo 进程在工作区 /blackboard/workspace 内完成 task.md 的任务, +# 用 --resume 保持跨轮连续上下文;完成时通过 blackboard 工具创建 DONE, +# supervisor 把结果以结构化 JSON(/blackboard/result.json)落盘。 +# +# 基础镜像复用本机已有的 pigo-worker:latest(Alpine,已含 bash/git/wget/flock/ +# timeout/CA 证书与静态 pigo 二进制),这里仅覆盖 pigo 二进制为新代码(含 +# blackboard 工具)并加入 supervisor 与协作 prompt,避免重新拉取 golang/alpine。 + +FROM pigo-worker:latest + +# solve/pentest tools: base image only ships busybox (no curl/python3), add them here. +# pigo-worker 镜像 USER=agent,RUN 层以非 root 执行,apk 无法锁库,先切 root 再恢复。 +USER root +RUN apk add --no-cache curl wget python3 py3-pip jq openssl + +# Python 常用库:网络/加密/数据/解析等,覆盖多数协作与 CTF 场景。 +# 使用清华 PyPI 镜像加速国内构建;--no-cache-dir 避免膨胀镜像。 +# --break-system-packages:Alpine 3.20+ 标记 externally-managed,pip 全局安装需显式放行。 +RUN pip3 install --no-cache-dir --break-system-packages \ + -i https://pypi.tuna.tsinghua.edu.cn/simple \ + requests urllib3 certifi idna charset-normalizer \ + cryptography pyOpenSSL paramiko pyjwt \ + PyYAML beautifulsoup4 lxml \ + python-dotenv click tqdm \ + numpy pandas \ + httpx aiohttp websockets dnspython \ + psutil pillow + +USER agent + +# 覆盖 pigo 二进制:coop/tmp/pigo-linux-amd64 由本地 Go 交叉编译(GOOS=linux CGO_ENABLED=0)。 +# legacy builder 不支持 COPY --chmod(需 BuildKit)。权限位由主机侧 chmod 提供(构建前先 chmod 755)。 +COPY coop/tmp/pigo-linux-amd64 /usr/local/bin/pigo +COPY coop/supervisor.sh /usr/local/bin/supervisor.sh +COPY coop/prompts /prompts +COPY coop/AGENTS.md /prompts/AGENTS.md + +WORKDIR /blackboard +ENTRYPOINT ["/usr/local/bin/supervisor.sh"] diff --git a/pigo/coop/README.md b/pigo/coop/README.md new file mode 100644 index 0000000..e978ea9 --- /dev/null +++ b/pigo/coop/README.md @@ -0,0 +1,196 @@ +# pigo 协作 —— 调用文档 + +> 在 Docker 容器内运行一个 pigo agent(单 agent)完成 /blackboard/task.md 中的任务,工作区为 /blackboard/workspace,完成时通过 blackboard 工具创建 DONE;supervisor 把结果以结构化 JSON(result.json)落盘,供外部调度方直接解析。 + +--- + +## 1. 原理与架构 + +``` + +------------------ docker 容器 pigo-coop ----------------+ + | | + TASK/MODEL/ | +---------- supervisor.sh(单 agent 编排)-------+ | + API_KEY ──────┼─▶ | 每轮运行一个 pigo 进程: | | + | | └─ agent ──▶ cwd = /blackboard/workspace | | + | | (--resume 会话 · 角色 prompt agent.md) | | + | +──────┬────────────────────────────────----------+ | + | ▼ | + | /blackboard(任务黑板:task.md + workspace + DONE) | + +-----------------------------------------------------------+ +``` + +- **单 agent**:每轮一个 pigo 进程串行推进,无评审者;工作区 `/blackboard/workspace` 是普通文件工具(read/write/bash)的根目录。 +- **原子完成标记**:`DONE` 只能通过 `blackboard` 工具 `done`(`O_CREATE|O_EXCL`)创建,supervisor 检测到 `DONE` 即收尾。 +- **连续上下文**:headless 运行自动持久化会话,supervisor 从 stream-json 首事件提取 `sessionId`,下一轮用 `--resume` 恢复,agent 跨轮记住自己的思考。 +- **协议注入**:任务协议(AGENTS.md)副本放入工作区,pigo 自动注入系统提示;角色 prompt 通过 `--append-system-prompt` 注入。 +- **结构化结果**:任务结束(成功/未解出/超时/失败)统一写 `$BB/result.json`,字段:`status`、`exit_code`、`summary`、`flag`、`artifacts`。 + +--- + +## 2. 前置条件 + +| 依赖 | 说明 | +| --- | --- | +| Go 工具链 | 版本与 `go.mod` 一致(当前 `go 1.27rc1`),用于交叉编译 pigo 二进制 | +| Docker | 本机已有基础镜像 `pigo-worker:latest`(Alpine,含 bash/git/wget/flock/timeout/CA 证书) | +| 模型后端 | 任意 OpenAI 兼容 API(如 DeepSeek、Ollama、OpenRouter),需可访问的 base-url 与 API key | + +--- + +## 3. 构建 + +### 3.1 交叉编译 pigo(含 blackboard 工具) + +Windows PowerShell: + +```powershell +$env:GOOS="linux"; $env:CGO_ENABLED="0" +go build -trimpath -ldflags="-s -w" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo +``` + +Linux/macOS: + +```bash +GOOS=linux CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo +``` + +产物 `coop/tmp/pigo-linux-amd64` 是纯静态 Linux 二进制(约 30MB)。 + +### 3.2 构建镜像 + +```bash +docker build -f coop/Dockerfile -t pigo-coop . +``` + +> 镜像直接 `FROM pigo-worker:latest`(复用本机已有镜像,不拉取 golang/alpine),只覆盖 pigo 二进制并加入 supervisor 与协作 prompt,构建通常秒级完成。 + +--- + +## 4. 环境变量(运行参数) + +| 变量 | 必填 | 默认值 | 说明 | +| --- | --- | --- | --- | +| `MODEL` | ✅ | — | 模型名,如 `deepseek-chat` | +| `BASE_URL` | ✅ | — | OpenAI 兼容 API base-url | +| `API_KEY` | ✅ | — | API key | +| `PROTOCOL` | 否 | 按 model 推断 | 协议 `openai` / `anthropic` | +| `TASK` | 二选一 | — | 任务描述;不填则需挂载 `$BB/task.md` | +| `BLACKBOARD` | 否 | `/blackboard` | 黑板根目录(即 `BB`) | +| `PROMPTS` | 否 | `/prompts` | 协作 prompt 目录(镜像内固定,一般不用改) | +| `ROUND_MAX` | 否 | `10` | 最大轮次 | +| `TIMEOUT` | 否 | `600` | 单轮超时秒数,`0`=不超时 | +| `FAIL_MODE` | 否 | `stop` | agent 失败时:`stop`=立即退出(默认)\| `continue`=继续下一轮 | + +--- + +## 5. 运行示例 + +### 5.1 最小运行(真实模型) + +```bash +docker run --rm \ + -e MODEL=deepseek-chat \ + -e BASE_URL=https://api.deepseek.com \ + -e API_KEY= \ + -e TASK="请为 XX 编写设计文档并实现原型" \ + -e ROUND_MAX=6 \ + pigo-coop +``` + +### 5.2 保留黑板产物与连续会话(推荐) + +```bash +mkdir -p blackboard +docker run --rm \ + -v $PWD/blackboard:/blackboard \ # 黑板产物落盘 + -v $HOME/.pigo:/root/.pigo \ # 保留 agent 会话(--resume 依赖) + -e MODEL=deepseek-chat \ + -e BASE_URL=https://api.deepseek.com \ + -e API_KEY= \ + -e TASK="..." \ + pigo-coop +``` + +> 会话持久化在 `~/.pigo/sessions`(容器内为 `/root/.pigo`)。**不挂载 `$HOME` 时,会话只存在于容器生命周期内**——跨 `docker run` 不保留;每次 `run` 内部跨轮仍有效。 + +### 5.3 结果与退出码 + +- 任务完成(检测到 `DONE`):退出码 `0`,`result.json` 的 `status=solved`。 +- 达到 `ROUND_MAX` 仍未完成:退出码 `1`,`status=unsolved`,检查黑板与日志。 +- 单轮超时被强杀:`status=timeout`(退出码 124)。 +- 参数缺失 / 非法:退出码 `2`。 + +--- + +## 6. 黑板目录结构 + +``` +/blackboard/ +├── task.md # 任务描述(勿修改) +├── AGENTS.md # 任务协议(工作区有副本) +├── workspace/ # agent 工作区 = 其 cwd(产物全部在这里) +├── DONE # 完成标记(只用 blackboard done 创建) +├── result.json # 结构化结果(supervisor 结束时生成) +├── sessions/ # agent 的 session id(supervisor 内部维护,供 --resume) +└── logs/ # 每轮运行日志:round-N.log +``` + +--- + +## 7. blackboard 工具(agent 侧 API) + +pigo 在 `BB` 环境变量存在时自动注册 `blackboard` 工具;普通运行不受影响。agent 通过该工具原子读写黑板: + +| 操作 | 参数 | 行为 | +| --- | --- | --- | +| `read` | (可选 `path`) | 无 `path`:返回黑板全局快照(task.md、workspace 清单、DONE 状态、环境);有 `path`:返回黑板内单文件内容(如 `workspace/exploit.py`) | +| `post` | `file`、`content` | 原子追加消息到 `messages/`;`file` 必须是裸 `*.md` 名(防目录穿越),内容 ≤ 32 KiB | +| `done` | `summary` | 原子创建 `DONE`(`O_CREATE|O_EXCL`);已存在则报错且不覆盖。仅交付完整时调用 | + +环境变量 `BB` / `ROUND` / `NAME` 由 supervisor 注入。 + +--- + +## 8. 验证与调试 + +### 8.1 无模型端到端验证(mock) + +`coop/tmp/mock_openai.py` 是一个 OpenAI 兼容 mock server,驱动 agent 依次调用 `blackboard done`,用于验证单 agent 编排与 DONE 原子性: + +```bash +# 终端 1:启动 mock(本机) +python coop/tmp/mock_openai.py 8899 + +# 终端 2:运行容器,base-url 指向 mock(按环境替换宿主地址) +docker run --rm \ + -v $PWD/blackboard:/blackboard \ + -e MODEL=mock -e BASE_URL=http://:8899 -e API_KEY=test \ + -e TASK="验证任务" -e ROUND_MAX=3 -e TIMEOUT=120 \ + pigo-coop +``` + +预期:agent 完成并创建 `DONE`,容器退出 0,`blackboard/result.json` 中 `status=solved`。 + +### 8.2 查看日志 + +```bash +# 黑板挂载后: +cat blackboard/logs/round-1.log # 第 1 轮的 stream-json 事件 +cat blackboard/sessions/agent.session # 下一轮 --resume 用的 session id +``` + +### 8.3 常见问题 + +| 现象 | 原因 / 处理 | +| --- | --- | +| 退出码 `2` 且提示缺 MODEL/BASE_URL/API_KEY | 未传必填环境变量 | +| 容器内报 `DONE already exists` | 正常:DONE 幂等,上一次运行已留标记 | +| 跨 run 后 agent 不记得之前轮次 | 未挂载 `$HOME/.pigo`;会话在容器内 `~/.pigo/sessions`,`--rm` 后丢失 | +| `docker build` 报找不到 `pigo-worker:latest` | 先 `docker pull` 或在目标机导入该基础镜像 | +| 构建期 `chmod` 被拒 | 已用 `COPY --chmod=755` 规避(NTFS 挂载构建上下文会丢可执行位) | + +--- + +## 9. 任务协议要点(详见 `/blackboard/AGENTS.md`) + +每轮:**读取 task.md → 检查已有产物 → 推进任务 → 完成后用 blackboard done 创建 DONE**。红线:不直接写 DONE;不伪造命令/提交结果;不修改 task.md 与 AGENTS.md。 diff --git a/pigo/coop/RUN.md b/pigo/coop/RUN.md new file mode 100644 index 0000000..a255f57 --- /dev/null +++ b/pigo/coop/RUN.md @@ -0,0 +1,107 @@ +# pigo 协作运行指南 + +> 只讲怎么跑。原理、架构、完整参数见 [README.md](./README.md)。 + +--- + +## 1. 前置检查(1 分钟) + +```bash +go version # 与 go.mod 一致(当前 go 1.27rc1) +docker images # 确认已有 pigo-worker:latest +``` + +没有 `pigo-worker:latest` 时先导入它,否则构建会失败。 + +--- + +## 2. 构建 + +```bash +# Windows PowerShell: +$env:GOOS="linux"; $env:CGO_ENABLED="0" +go build -trimpath -ldflags="-s -w" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo + +# Linux/macOS: +GOOS=linux CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o coop/tmp/pigo-linux-amd64 ./cmd/pigo +``` + +```bash +docker build -f coop/Dockerfile -t pigo-coop . +``` + +--- + +## 3. 运行 + +### 3.1 直接运行(不保留任何数据) + +```bash +docker run --rm \ + -e MODEL=deepseek-chat \ + -e BASE_URL=https://api.deepseek.com \ + -e API_KEY= \ + -e TASK="请为 XX 编写设计文档并实现原型" \ + -e ROUND_MAX=6 \ + pigo-coop +``` + +### 3.2 推荐运行(保留黑板产物 + 跨 run 会话) + +```bash +mkdir -p blackboard + +docker run --rm \ + -v $PWD/blackboard:/blackboard \ + -v $HOME/.pigo:/root/.pigo \ + -e MODEL=deepseek-chat \ + -e BASE_URL=https://api.deepseek.com \ + -e API_KEY= \ + -e TASK="请为 XX 编写设计文档并实现原型" \ + -e ROUND_MAX=6 \ + pigo-coop +``` + +> 不挂载 `$HOME/.pigo` 时,agent 会话只在容器生命周期内有效(单次 run 内跨轮不受影响)。 + +### 3.3 常用参数速查 + +| 参数 | 必填 | 默认 | 用途 | +| --- | --- | --- | --- | +| `MODEL` | ✅ | — | 模型名 | +| `BASE_URL` | ✅ | — | OpenAI 兼容 base-url | +| `API_KEY` | ✅ | — | API key | +| `TASK` | ✅(或用 `task.md`) | — | 任务描述 | +| `ROUND_MAX` | 否 | `10` | 最大轮次 | +| `TIMEOUT` | 否 | `600` | 单轮超时秒数(`0`=不限) | +| `PROTOCOL` | 否 | 按模型推断 | `openai` / `anthropic` | + +--- + +## 4. 结果怎么看 + +- **完成**:stdout 打印 `[supervisor] status=solved`,退出码 `0`。 +- **未完成**:达到 `ROUND_MAX`,`status=unsolved`,退出码 `1`。 +- **参数错**:退出码 `2`。 + +保留黑板后查看: + +```bash +ls blackboard/workspace/ # agent 产物 +cat blackboard/DONE # 最终交付总结 +cat blackboard/result.json # 结构化结果(status/exit_code/summary/flag/artifacts) +cat blackboard/logs/round-1.log # 第 1 轮运行日志 +``` + +`result.json` 是给外部调度方(如 agent-web 的 run_coop 工具)直接解析的标准结果文件,字段固定:`status`、`exit_code`、`summary`、`flag`、`artifacts`。 + +--- + +## 5. 常见运行问题 + +| 报错 / 现象 | 处理 | +| --- | --- | +| 退出 2:`必须提供 MODEL / BASE_URL / API_KEY` | 补传环境变量 | +| `DONE already exists` | 正常,上一次运行已留标记(幂等) | +| 找不到 `pigo-worker:latest` | 先导入基础镜像再 build | +| 想从头重跑 | 删除挂载目录里的 `DONE` 与 `result.json` 后重新 run | diff --git a/pigo/coop/build.ps1 b/pigo/coop/build.ps1 new file mode 100644 index 0000000..e5af4e1 --- /dev/null +++ b/pigo/coop/build.ps1 @@ -0,0 +1,19 @@ +# 构建 pigo 协作镜像(单 agent 版:一个 pigo 进程完成任务,blackboard 工具负责原子创建 DONE) +# 用法(仓库根目录执行): powershell -File coop/build.ps1 +$ErrorActionPreference = "Stop" + +docker build -f coop/Dockerfile -t pigo-coop . + +Write-Host "" +Write-Host "构建完成: pigo-coop" +Write-Host "" +Write-Host "运行示例(外部调用方传入模型/API 信息与任务):" +Write-Host ' docker run --rm \' +Write-Host ' -e MODEL=deepseek-chat \' +Write-Host ' -e BASE_URL=https://api.deepseek.com \' +Write-Host ' -e API_KEY= \' +Write-Host ' -e TASK="请为 XX 编写设计文档并实现原型" \' +Write-Host ' -e ROUND_MAX=6 \' +Write-Host ' pigo-coop' +Write-Host "" +Write-Host "黑板产物默认在容器内 /blackboard,如需保留可挂载卷: -v $PWD/blackboard:/blackboard" diff --git a/pigo/coop/prompts/agent.md b/pigo/coop/prompts/agent.md new file mode 100644 index 0000000..db6cb6f --- /dev/null +++ b/pigo/coop/prompts/agent.md @@ -0,0 +1,18 @@ +# 角色:agent(主执行者) + +你是 pigo 单 agent 任务执行者,独立完成 /blackboard/task.md 中的任务。 + +## 你的职责 +- 理解任务、制定方案、完成核心产出。 +- 所有产物写入你的工作区(你的 cwd,即 /blackboard/workspace/),按需命名。 +- 需要提交 flag 的任务:拿到 flag 后立即按任务中的提交规则提交,并把提交响应与得分记录在产物中。 + +## 工作方式 +- 每轮:读取 /blackboard/task.md 与工作区已有产物 → 推进任务 → 将进展与产物写入工作区。 +- 环境变量 ROUND、NAME、BB 由 supervisor 提供;需要记录进展或创建完成标记时用 blackboard 工具(BB 指向黑板根目录)。 +- 不要伪造命令结果、文件内容或提交响应;如实记录。 +- 全部工作完成、交付物完整时,用 blackboard 工具 action=done 创建完成标记(summary 写最终交付总结,含关键结果、flag、提交响应、产物清单)。 + +## 注意 +- 若本轮未完成,下一轮会用 --resume 恢复你的会话继续推进,跨轮保持上下文。 +- 宁可多轮打磨,不要过早宣布完成;确认无法解出时同样用 action=done 标记,summary 写明「未解出」与已尝试内容,让调度方及时切换下一题。 diff --git a/pigo/coop/supervisor.sh b/pigo/coop/supervisor.sh new file mode 100644 index 0000000..889a109 --- /dev/null +++ b/pigo/coop/supervisor.sh @@ -0,0 +1,182 @@ +#!/usr/bin/env bash +# pigo 协作 supervisor(单 agent 版) +# +# 在容器内运行一个 pigo headless 进程完成 /blackboard/task.md 中的任务。 +# agent 的工作区为 /blackboard/workspace,产物与轮次日志保留在黑板上; +# 任务完成(agent 用 blackboard 工具创建 DONE)后,supervisor 把结果 +# 以结构化 JSON(result.json)落盘并打印一行摘要,供外部调度方直接解析。 +# +# 与旧双 agent 版的关键差异: +# 1. 单 agent:每轮只运行一个 pigo 进程,无评审者,串行推进直至 DONE。 +# 2. 结构化结果:无论成功/未解出/超时/失败,统一写 $BB/result.json +# (status / exit_code / summary / flag / artifacts),主 agent 无需再 +# 漫游黑板文件逐段猜测结果。 +# 3. 日志精简:supervisor 只输出带 [supervisor] 前缀的状态行;agent 的 +# stream-json 事件只写日志文件,不再混入 stdout。 +# +# 环境变量(均由外部调用方传入): +# MODEL 模型名,如 deepseek-chat (必填) +# BASE_URL OpenAI 兼容 API base-url (必填) +# API_KEY API key (必填) +# PROTOCOL 协议 openai|anthropic (可选,默认按 model 推断) +# TASK 任务描述 (必填;或挂载 $BB/task.md) +# BLACKBOARD 黑板目录 (默认 /blackboard) +# PROMPTS 协作 prompt 目录 (默认 /prompts) +# ROUND_MAX 最大轮次 (默认 10) +# TIMEOUT 单轮超时秒数,0=不超时 (默认 600) +# FAIL_MODE agent 失败时:stop=立即退出(默认)| continue=继续下一轮 +set -u + +BB="${BLACKBOARD:-/blackboard}" +PROMPTS="${PROMPTS:-/prompts}" +ROUND_MAX="${ROUND_MAX:-10}" +TIMEOUT="${TIMEOUT:-600}" +FAIL_MODE="${FAIL_MODE:-stop}" + +log() { echo "[supervisor] $*"; } + +# emit_result 把任务结果以结构化 JSON 写入 $BB/result.json 并打印一行摘要。 +# status: solved | unsolved | error | timeout +emit_result() { + local status="$1" summary="$2" code="$3" + local flag="" artifacts_json="[]" summary_json='""' flag_json='""' + + # flag 优先从 summary 提取,其次在 workspace 产物中全量检索 + if [ -n "$summary" ]; then + flag=$(printf '%s' "$summary" | grep -oE 'flag\{[^}]+\}' | head -1) + fi + if [ -z "$flag" ] && [ -d "$BB/workspace" ]; then + flag=$(grep -rhoE 'flag\{[^}]+\}' "$BB/workspace" 2>/dev/null | head -1) + fi + + # 产物清单:workspace 下的全部文件(最多 50 个),经 python3 转义为 JSON 数组 + if [ -d "$BB/workspace" ]; then + artifacts_json=$( + cd "$BB/workspace" && find . -type f 2>/dev/null | sed 's|^\./||' | head -50 | + python3 -c 'import json,sys; print(json.dumps([l.rstrip("\n") for l in sys.stdin]))' 2>/dev/null + ) + [ -z "$artifacts_json" ] && artifacts_json="[]" + fi + + # summary 截断到 4000 字符并经 python3 转义,防止引号/换行破坏 JSON + summary_json=$(printf '%s' "$summary" | head -c 4000 | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))' 2>/dev/null) + [ -z "$summary_json" ] && summary_json='""' + flag_json=$(printf '%s' "$flag" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))' 2>/dev/null) + [ -z "$flag_json" ] && flag_json='""' + + cat > "$BB/result.json" <&2 + exit 2 +fi +if [ -z "${TASK:-}" ] && [ ! -f "$BB/task.md" ]; then + echo "错误:请通过 TASK 环境变量或挂载 $BB/task.md 提供任务" >&2 + exit 2 +fi +if [ -f "$BB/DONE" ]; then + log "黑板已有完成标记,如需重新开始请删除 $BB/DONE" + exit 0 +fi + +# ---- 初始化黑板 ---- +mkdir -p "$BB/workspace" "$BB/logs" "$BB/sessions" +if [ -n "${TASK:-}" ]; then + printf '%s\n' "$TASK" > "$BB/task.md" +fi +# 协议注入:agent 的 cwd 是 $BB/workspace,AGENTS.md 注入链从 cwd 起, +# 因此把协议副本放进工作区(勿修改,它是每轮系统提示的一部分) +cp "$PROMPTS/AGENTS.md" "$BB/workspace/AGENTS.md" +log "黑板初始化完成:$BB" +log "任务:$(head -c 200 "$BB/task.md")" + +# ---- 单轮运行:一个 pigo 进程,输出 stream-json(首事件携带 sessionId)供 --resume ---- +run_agent() { + local round="$1" + local session_file="$BB/sessions/agent.session" + local task_prompt="轮次 $round 开始。先读取 /blackboard/task.md 中的任务,检查工作区已有产物与 DONE 状态,然后继续推进任务。全部工作完成、交付物完整时,用 blackboard 工具 action=done 创建完成标记(summary 写最终交付总结,含关键结果/flag/提交响应/产物清单)。" + local args=(-p "$task_prompt" -a -C "$BB/workspace" -o stream-json \ + --model "$MODEL" --base-url "$BASE_URL" --api-key "$API_KEY" \ + --append-system-prompt "$PROMPTS/agent.md") + [ -n "${PROTOCOL:-}" ] && args+=(--protocol "$PROTOCOL") + if [ -s "$session_file" ]; then + args+=(--resume "$(cat "$session_file")") + fi + + export ROUND="$round" NAME="agent" BB="$BB" + local logfile="$BB/logs/round-$round.log" + log "第 $round 轮开始运行 agent(日志:$logfile)" + if [ "$TIMEOUT" -gt 0 ] 2>/dev/null; then + timeout "$TIMEOUT" pigo "${args[@]}" > "$logfile" 2>&1 + else + pigo "${args[@]}" > "$logfile" 2>&1 + fi + local rc=$? + + # 从 stream-json 首事件提取 sessionId 供下一轮 --resume(agent_start 事件带 env.sessionId) + local sid + sid=$(grep -o '"sessionId":"[^"]*"' "$logfile" 2>/dev/null | head -1 | cut -d'"' -f4) + if [ -n "$sid" ]; then + printf '%s' "$sid" > "$session_file" + fi + return $rc +} + +# ---- 主循环:逐轮运行 agent,直到 DONE 或达到轮次上限 ---- +for round in $(seq 1 "$ROUND_MAX"); do + [ -f "$BB/DONE" ] && break + log "===== 第 $round 轮开始 =====" + + run_agent "$round" + rc=$? + + if [ "$rc" -eq 0 ]; then + log "第 $round 轮完成" + else + log "第 $round 轮失败(exit=$rc),日志见 $BB/logs/round-$round.log" + if [ "$FAIL_MODE" = "stop" ]; then + # 兜底:agent 虽超时/失败,但黑板已有提交成功证据(workspace 内 *.md 含 correct:true) + # → 视为完成正常退出,避免"flag 已提交却因未写 DONE 被强杀(exit 143)"。 + _done=0 + for _f in "$BB"/workspace/*.md "$BB"/messages/*.md; do + [ -f "$_f" ] && grep -q '"correct":true' "$_f" && _done=1 && break + done + if [ "$_done" = "1" ]; then + log "检测到提交成功证据,自动标记完成" + { echo "Task finished (auto-detected submit success)"; } > "$BB/DONE" + emit_result solved "$(cat "$BB/DONE")" 0 + exit 0 + fi + if [ "$rc" -eq 124 ]; then + emit_result timeout "" "$rc" + else + emit_result error "agent 退出码 $rc,日志见 logs/round-$round.log" "$rc" + fi + exit "$rc" + fi + fi +done + +if [ -f "$BB/DONE" ]; then + log "检测到完成标记,任务结束" + emit_result solved "$(cat "$BB/DONE")" 0 + exit 0 +fi + +log "达到最大轮次 $ROUND_MAX 仍未完成,请检查黑板产物与日志" +emit_result unsolved "" 1 +exit 1 diff --git a/pigo/go.mod b/pigo/go.mod new file mode 100644 index 0000000..bb36da1 --- /dev/null +++ b/pigo/go.mod @@ -0,0 +1,68 @@ +module github.com/smallnest/pigo + +go 1.27rc1 + +require ( + charm.land/bubbles/v2 v2.1.1 + charm.land/bubbletea/v2 v2.0.8 + charm.land/lipgloss/v2 v2.0.5 + github.com/BurntSushi/toml v1.6.0 + github.com/JohannesKaufmann/html-to-markdown/v2 v2.5.2 + github.com/charmbracelet/glamour v1.0.0 + github.com/charmbracelet/x/ansi v0.11.7 + github.com/coder/websocket v1.8.13 + github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 + github.com/openai/openai-go v1.12.0 + github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 + github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e + github.com/spf13/pflag v1.0.10 + golang.org/x/net v0.57.0 + golang.org/x/text v0.40.0 + gopkg.in/yaml.v3 v3.0.1 + modernc.org/sqlite v1.55.0 +) + +require ( + github.com/JohannesKaufmann/dom v0.3.1 // indirect + github.com/alecthomas/chroma/v2 v2.20.0 // indirect + github.com/atotto/clipboard v0.1.4 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect + github.com/aymerick/douceur v0.2.0 // indirect + github.com/charmbracelet/colorprofile v0.4.3 // indirect + github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 // indirect + github.com/charmbracelet/ultraviolet v0.0.0-20260703014108-f5a850f9c2b7 // indirect + github.com/charmbracelet/x/cellbuf v0.0.15 // indirect + github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf // indirect + github.com/charmbracelet/x/term v0.2.2 // indirect + github.com/charmbracelet/x/termios v0.1.1 // indirect + github.com/charmbracelet/x/windows v0.2.2 // indirect + github.com/clipperhouse/displaywidth v0.11.0 // indirect + github.com/clipperhouse/uax29/v2 v2.7.0 // indirect + github.com/dlclark/regexp2 v1.11.5 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/gorilla/css v1.0.1 // indirect + github.com/lucasb-eyer/go-colorful v1.4.0 // indirect + github.com/mattn/go-isatty v0.0.22 // indirect + github.com/mattn/go-runewidth v0.0.24 // indirect + github.com/microcosm-cc/bluemonday v1.0.27 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/muesli/reflow v0.3.0 // indirect + github.com/muesli/termenv v0.16.0 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/rivo/uniseg v0.4.7 // indirect + github.com/tidwall/gjson v1.14.4 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect + github.com/yuin/goldmark v1.8.2 // indirect + github.com/yuin/goldmark-emoji v1.0.6 // indirect + golang.org/x/sync v0.22.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/term v0.45.0 // indirect + modernc.org/libc v1.74.1 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect +) diff --git a/pigo/go.sum b/pigo/go.sum new file mode 100644 index 0000000..6111626 --- /dev/null +++ b/pigo/go.sum @@ -0,0 +1,173 @@ +charm.land/bubbles/v2 v2.1.1 h1:7r55WzBxpo/R3z98hGmY7KKPd3ET6vsf0Fb9sDHOV60= +charm.land/bubbles/v2 v2.1.1/go.mod h1:GE6M31gaWZVXzGw73OeuTTgy4lX+OtkH0E5ymnNsHxo= +charm.land/bubbletea/v2 v2.0.8 h1:SxTJMhCAI3lbPmy4SgX5LWZ24AdINr4I6UEqzZvYJuY= +charm.land/bubbletea/v2 v2.0.8/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss= +charm.land/lipgloss/v2 v2.0.5 h1:kbNxgeeUOYv5J0YdpxFjfvf3dFvqH8Aci4zB6xqFtrY= +charm.land/lipgloss/v2 v2.0.5/go.mod h1:9oqhxt4yxIMe6q5A4kHr44DremZk7J9UNh74GlWa5nc= +github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk= +github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= +github.com/JohannesKaufmann/dom v0.3.1 h1:J16l9JAHWgkFPR3VIPbQ1gvS0cWab6laK1q7PFL3qh0= +github.com/JohannesKaufmann/dom v0.3.1/go.mod h1:BZPkf8ZeYrBgABjwJn9iiKt8aiCtkxpHkevms+Yp2DE= +github.com/JohannesKaufmann/html-to-markdown/v2 v2.5.2 h1:XFJZFWESIWlUEHHjzBuv8RvrtCWnSGlimEX17ysSDb8= +github.com/JohannesKaufmann/html-to-markdown/v2 v2.5.2/go.mod h1:BHWO8lJzttJLqwuV8Rb1B3OG2OSzLbssZDI1FRg2eAA= +github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= +github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= +github.com/alecthomas/assert/v2 v2.11.0 h1:2Q9r3ki8+JYXvGsDyBXwH3LcJ+WK5D0gc5E8vS6K3D0= +github.com/alecthomas/assert/v2 v2.11.0/go.mod h1:Bze95FyfUr7x34QZrjL+XP+0qgp/zg8yS+TtBj1WA3k= +github.com/alecthomas/chroma/v2 v2.20.0 h1:sfIHpxPyR07/Oylvmcai3X/exDlE8+FA820NTz+9sGw= +github.com/alecthomas/chroma/v2 v2.20.0/go.mod h1:e7tViK0xh/Nf4BYHl00ycY6rV7b8iXBksI9E359yNmA= +github.com/alecthomas/repr v0.5.1 h1:E3G4t2QbHTSNpPKBgMTln5KLkZHLOcU7r37J4pXBuIg= +github.com/alecthomas/repr v0.5.1/go.mod h1:Fr0507jx4eOXV7AlPV6AVZLYrLIuIeSOWtW57eE/O/4= +github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= +github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= +github.com/aymanbagabas/go-udiff v0.4.1 h1:OEIrQ8maEeDBXQDoGCbbTTXYJMYRCRO1fnodZ12Gv5o= +github.com/aymanbagabas/go-udiff v0.4.1/go.mod h1:0L9PGwj20lrtmEMeyw4WKJ/TMyDtvAoK9bf2u/mNo3w= +github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= +github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= +github.com/charmbracelet/colorprofile v0.4.3 h1:QPa1IWkYI+AOB+fE+mg/5/4HRMZcaXex9t5KX76i20Q= +github.com/charmbracelet/colorprofile v0.4.3/go.mod h1:/zT4BhpD5aGFpqQQqw7a+VtHCzu+zrQtt1zhMt9mR4Q= +github.com/charmbracelet/glamour v1.0.0 h1:AWMLOVFHTsysl4WV8T8QgkQ0s/ZNZo7CiE4WKhk8l08= +github.com/charmbracelet/glamour v1.0.0/go.mod h1:DSdohgOBkMr2ZQNhw4LZxSGpx3SvpeujNoXrQyH2hxo= +github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 h1:ZR7e0ro+SZZiIZD7msJyA+NjkCNNavuiPBLgerbOziE= +github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834/go.mod h1:aKC/t2arECF6rNOnaKaVU6y4t4ZeHQzqfxedE/VkVhA= +github.com/charmbracelet/ultraviolet v0.0.0-20260703014108-f5a850f9c2b7 h1:3FmWoGNWK4STvqg0O0Aeav2T7rodWJAPeF0QpH+8gFw= +github.com/charmbracelet/ultraviolet v0.0.0-20260703014108-f5a850f9c2b7/go.mod h1:f/jRa757WUmaOZrbPspXymbg/GnbF+rwe4OLsG7aXYo= +github.com/charmbracelet/x/ansi v0.11.7 h1:kzv1kJvjg2S3r9KHo8hDdHFQLEqn4RBCb39dAYC84jI= +github.com/charmbracelet/x/ansi v0.11.7/go.mod h1:9qGpnAVYz+8ACONkZBUWPtL7lulP9No6p1epAihUZwQ= +github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI= +github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q= +github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f h1:pk6gmGpCE7F3FcjaOEKYriCvpmIN4+6OS/RD0vm4uIA= +github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f/go.mod h1:IfZAMTHB6XkZSeXUqriemErjAWCCzT0LwjKFYCZyw0I= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf h1:rLG0Yb6MQSDKdB52aGX55JT1oi0P0Kuaj7wi1bLUpnI= +github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf/go.mod h1:B3UgsnsBZS/eX42BlaNiJkD1pPOUa+oF1IYC6Yd2CEU= +github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= +github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= +github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY= +github.com/charmbracelet/x/termios v0.1.1/go.mod h1:rB7fnv1TgOPOyyKRJ9o+AsTU/vK5WHJ2ivHeut/Pcwo= +github.com/charmbracelet/x/windows v0.2.2 h1:IofanmuvaxnKHuV04sC0eBy/smG6kIKrWG2/jYn2GuM= +github.com/charmbracelet/x/windows v0.2.2/go.mod h1:/8XtdKZzedat74NQFn0NGlGL4soHB0YQZrETF96h75k= +github.com/clipperhouse/displaywidth v0.11.0 h1:lBc6kY44VFw+TDx4I8opi/EtL9m20WSEFgwIwO+UVM8= +github.com/clipperhouse/displaywidth v0.11.0/go.mod h1:bkrFNkf81G8HyVqmKGxsPufD3JhNl3dSqnGhOoSD/o0= +github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJv2v7Vk= +github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= +github.com/coder/websocket v1.8.13 h1:f3QZdXy7uGVz+4uCJy2nTZyM0yTBj8yANEHhqlXZ9FE= +github.com/coder/websocket v1.8.13/go.mod h1:LNVeNrXQZfe5qhS9ALED3uA+l5pPqvwXg3CKoDBB2gs= +github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ= +github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8= +github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= +github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= +github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs= +github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8= +github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= +github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= +github.com/mattn/go-isatty v0.0.22 h1:j8l17JJ9i6VGPUFUYoTUKPSgKe/83EYU2zBC7YNKMw4= +github.com/mattn/go-isatty v0.0.22/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4= +github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk= +github.com/mattn/go-runewidth v0.0.24 h1:cpokDiIn0MGnhdHwuWnJBITySJ20QyNGnY2kR/ay2DU= +github.com/mattn/go-runewidth v0.0.24/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= +github.com/microcosm-cc/bluemonday v1.0.27 h1:MpEUotklkwCSLeH+Qdx1VJgNqLlpY2KXwXFM08ygZfk= +github.com/microcosm-cc/bluemonday v1.0.27/go.mod h1:jFi9vgW+H7c3V0lb6nR74Ib/DIB5OBs92Dimizgw2cA= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= +github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= +github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= +github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0= +github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ= +github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= +github.com/sebdah/goldie/v2 v2.8.0 h1:dZb9wR8q5++oplmEiJT+U/5KyotVD+HNGCAc5gNr8rc= +github.com/sebdah/goldie/v2 v2.8.0/go.mod h1:oZ9fp0+se1eapSRjfYbsV/0Hqhbuu3bJVvKI/NNtssI= +github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw= +github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4= +github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e h1:MRM5ITcdelLK2j1vwZ3Je0FKVCfqOLp5zO6trqMLYs0= +github.com/skip2/go-qrcode v0.0.0-20200617195104-da1b6568686e/go.mod h1:XV66xRDqSt+GTGFMVlhk3ULuV0y9ZmzeVGR4mloJI3M= +github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk= +github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM= +github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= +github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= +github.com/yuin/goldmark v1.8.2 h1:kEGpgqJXdgbkhcOgBxkC0X0PmoPG1ZyoZ117rDVp4zE= +github.com/yuin/goldmark v1.8.2/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg= +github.com/yuin/goldmark-emoji v1.0.6 h1:QWfF2FYaXwL74tfGOW5izeiZepUDroDJfWubQI9HTHs= +github.com/yuin/goldmark-emoji v1.0.6/go.mod h1:ukxJDKFpdFb5x0a5HqbdlcKtebh086iJpI31LTKmWuA= +golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI= +golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo= +golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= +golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= +golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= +golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0= +golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= +golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= +golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.29.0 h1:CXgwL8cvxmyzBQZzbSl/6xFtMCryb6u8IOqDci39cgc= +modernc.org/cc/v4 v4.29.0/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.6 h1:sBgfIwyN0TQ9C5hwIeuqyeAKyMWnbvj2fvpF4L11uzU= +modernc.org/ccgo/v4 v4.34.6/go.mod h1:SZ8YcN9NG7XVsQYdm6jYBvi8PQP1qi+kqB6OhjqI3Fk= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.4 h1:2g65LGVSmFQrXeITAw97x7hCRvZFcyE1uDP+7Vng7JI= +modernc.org/gc/v3 v3.1.4/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.74.1 h1:bdR4VTKFMC4966QSNZ05XLGI/VwzVa2kTUX51Dm0riQ= +modernc.org/libc v1.74.1/go.mod h1:uH4t5bOx3G3g9Xcmj10YKlTcVISlRDwv8VoQJG9n8Os= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.55.0 h1:hIFh0MCH0rGinQ/4KYb5/UbCkRkb+UP+OkLCVWa5MTM= +modernc.org/sqlite v1.55.0/go.mod h1:4ntCLuNmnH8+GNqjka1wNg7KJd5/Hi5FYp8K+XQ7GZw= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/pigo/install.sh b/pigo/install.sh new file mode 100644 index 0000000..e77ce10 --- /dev/null +++ b/pigo/install.sh @@ -0,0 +1,117 @@ +#!/bin/sh +# pigo 安装脚本:检测当前操作系统 / 架构,从 GitHub Releases 下载最新的 +# 预编译二进制,并安装到常用的 PATH 目录。 +# +# 用法: +# curl -fsSL https://raw.githubusercontent.com/smallnest/pigo/master/install.sh | sh +# +# 可用环境变量覆盖默认行为: +# PIGO_VERSION 指定版本(形如 v0.2.0),默认取最新 release +# PIGO_INSTALL_DIR 安装目录,默认 /usr/local/bin(无写权限时回退到 ~/.local/bin) +# GITHUB_TOKEN 可选,用于提高 GitHub API 速率限制 +set -eu + +REPO="smallnest/pigo" +BINARY="pigo" + +info() { printf '%s\n' "pigo-install: $*" >&2; } +err() { printf '%s\n' "pigo-install: error: $*" >&2; exit 1; } + +need() { command -v "$1" >/dev/null 2>&1 || err "缺少依赖命令: $1"; } + +# 1. 检测下载器(curl 或 wget)。 +if command -v curl >/dev/null 2>&1; then + DL="curl -fsSL" + DLO="curl -fsSL -o" +elif command -v wget >/dev/null 2>&1; then + DL="wget -qO-" + DLO="wget -qO" +else + err "需要 curl 或 wget" +fi +need tar +need uname + +# 2. 检测 OS,映射到 goreleaser 的归档命名(见 .goreleaser.yaml)。 +os_raw=$(uname -s) +case "$os_raw" in + Linux) OS="Linux" ;; + Darwin) OS="Darwin" ;; + MINGW* | MSYS* | CYGWIN* | Windows_NT) + err "Windows 请从 Releases 页面下载 .zip:https://github.com/$REPO/releases" ;; + *) err "不支持的操作系统: $os_raw" ;; +esac + +# 3. 检测架构,映射到归档命名(amd64→x86_64,386→i386,arm64 保持)。 +arch_raw=$(uname -m) +case "$arch_raw" in + x86_64 | amd64) ARCH="x86_64" ;; + arm64 | aarch64) ARCH="arm64" ;; + i386 | i686) ARCH="i386" ;; + *) err "不支持的架构: $arch_raw" ;; +esac + +# 4. 解析目标版本:优先 PIGO_VERSION,否则查询最新 release 的 tag。 +VERSION="${PIGO_VERSION:-}" +api_auth="" +[ -n "${GITHUB_TOKEN:-}" ] && api_auth="-H Authorization:\ Bearer\ $GITHUB_TOKEN" +if [ -z "$VERSION" ]; then + info "查询最新 release ..." + # 从 GitHub API 的 latest 端点提取 tag_name。 + latest_json=$($DL "https://api.github.com/repos/$REPO/releases/latest" 2>/dev/null) || \ + err "无法访问 GitHub API,请检查网络或用 PIGO_VERSION 指定版本" + VERSION=$(printf '%s' "$latest_json" | grep -o '"tag_name"[ ]*:[ ]*"[^"]*"' | head -n1 | sed 's/.*"tag_name"[ ]*:[ ]*"\([^"]*\)".*/\1/') + [ -n "$VERSION" ] || err "无法解析最新版本号,请用 PIGO_VERSION 指定" +fi + +# 归档名里的版本号不带前导 v(goreleaser 的 .Version)。 +VER_NUM=$(printf '%s' "$VERSION" | sed 's/^v//') +ARCHIVE="${BINARY}_${VER_NUM}_${OS}_${ARCH}.tar.gz" +URL="https://github.com/$REPO/releases/download/$VERSION/$ARCHIVE" + +info "版本: $VERSION" +info "平台: ${OS}/${ARCH}" +info "下载: $URL" + +# 5. 下载并解压到临时目录。 +TMP=$(mktemp -d 2>/dev/null || mktemp -d -t pigo-install) +trap 'rm -rf "$TMP"' EXIT INT TERM +$DLO "$TMP/$ARCHIVE" "$URL" || err "下载失败: $URL" +tar -xzf "$TMP/$ARCHIVE" -C "$TMP" || err "解压失败: $ARCHIVE" +[ -f "$TMP/$BINARY" ] || err "归档中未找到二进制 $BINARY" +chmod +x "$TMP/$BINARY" + +# 6. 选择安装目录:PIGO_INSTALL_DIR > /usr/local/bin > ~/.local/bin。 +DIR="${PIGO_INSTALL_DIR:-}" +if [ -z "$DIR" ]; then + if [ -w /usr/local/bin ] 2>/dev/null; then + DIR="/usr/local/bin" + elif [ "$(id -u)" = "0" ]; then + DIR="/usr/local/bin" + else + DIR="$HOME/.local/bin" + fi +fi +mkdir -p "$DIR" 2>/dev/null || err "无法创建安装目录: $DIR" + +# 7. 安装。若目录不可写但可 sudo,尝试用 sudo。 +DEST="$DIR/$BINARY" +if [ -w "$DIR" ]; then + mv "$TMP/$BINARY" "$DEST" +elif command -v sudo >/dev/null 2>&1; then + info "$DIR 需要提升权限,使用 sudo 安装 ..." + sudo mv "$TMP/$BINARY" "$DEST" +else + err "$DIR 不可写且无 sudo,请设置 PIGO_INSTALL_DIR 指向可写目录" +fi + +info "已安装: $DEST" + +# 8. 提示 PATH 是否包含安装目录。 +case ":$PATH:" in + *":$DIR:"*) : ;; + *) info "注意: $DIR 不在 PATH 中,请将其加入 PATH,例如:" >&2 + info " echo 'export PATH=\"$DIR:\$PATH\"' >> ~/.profile" >&2 ;; +esac + +"$DEST" --version || true diff --git a/pigo/internal/agentcore/content.go b/pigo/internal/agentcore/content.go new file mode 100644 index 0000000..48a0036 --- /dev/null +++ b/pigo/internal/agentcore/content.go @@ -0,0 +1,195 @@ +// Package agentcore defines the core "leaf" data types and control flow for the +// pigo agent harness, a Go reimplementation of the pi agent loop. It is the +// foundation package that every other agent sub-package depends on and imports +// nothing from them. +// +// This file defines the Content model: a sealed interface implemented by the +// four content block kinds (text, thinking, toolCall, image). Because Go's +// encoding/json cannot dispatch to an interface based on a discriminant field, +// containers holding []Content implement custom UnmarshalJSON that peeks at the +// "type" field and decodes into the concrete struct. Mirrors pi's discriminated +// union (packages/ai/src/types.ts) as interface + type switch. +package agentcore + +import ( + "bytes" + "encoding/json" + "fmt" +) + +// Content is a sealed interface implemented by every content block kind. +// Consumers dispatch with a type switch. The interface is sealed via the +// unexported isContent marker so no type outside this package can satisfy it. +type Content interface { + isContent() +} + +// Content type discriminants, matching pi's wire format. +const ( + ContentTypeText = "text" + ContentTypeThinking = "thinking" + ContentTypeToolCall = "toolCall" + ContentTypeImage = "image" +) + +// TextContent is a plain text block. +type TextContent struct { + Type string `json:"type"` + Text string `json:"text"` + TextSignature string `json:"textSignature,omitempty"` +} + +// ThinkingContent is a reasoning/thinking block. Never folded into text. +type ThinkingContent struct { + Type string `json:"type"` + Thinking string `json:"thinking"` + ThinkingSignature string `json:"thinkingSignature,omitempty"` + Redacted bool `json:"redacted,omitempty"` +} + +// ToolCallContent is a request from the model to invoke a tool. Arguments are +// kept as raw JSON so validation (JSON Schema) and shaping happen downstream. +type ToolCallContent struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + ThoughtSignature string `json:"thoughtSignature,omitempty"` +} + +// ImageContent is an image block (base64 data + mime type). +type ImageContent struct { + Type string `json:"type"` + Data string `json:"data"` + MimeType string `json:"mimeType"` +} + +func (TextContent) isContent() {} +func (ThinkingContent) isContent() {} +func (ToolCallContent) isContent() {} +func (ImageContent) isContent() {} + +// MarshalJSON encodes a ToolCallContent, tolerating malformed Arguments. A +// model can stream syntactically invalid tool-call JSON (a truncated or +// duplicated key, e.g. `{"todos": []{}...`); such bytes are kept verbatim in +// Arguments so schema validation can report "not valid JSON" to the model, but +// json.RawMessage.MarshalJSON rejects them, which would otherwise abort every +// downstream serialization (session persistence, provider re-serialization) and +// take the whole turn down. To keep those paths crash-free we emit invalid +// arguments as a JSON string of the raw bytes: valid JSON that round-trips the +// original text. Well-formed arguments are emitted unchanged. +func (t ToolCallContent) MarshalJSON() ([]byte, error) { + args := t.Arguments + if len(bytes.TrimSpace(args)) == 0 { + args = json.RawMessage("{}") + } else if !json.Valid(args) { + s, err := json.Marshal(string(args)) + if err != nil { + return nil, fmt.Errorf("content: encode invalid tool arguments: %w", err) + } + args = s + } + // A named alias avoids recursing into this MarshalJSON. + type wire struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + ThoughtSignature string `json:"thoughtSignature,omitempty"` + } + return json.Marshal(wire{ + Type: t.Type, + ID: t.ID, + Name: t.Name, + Arguments: args, + ThoughtSignature: t.ThoughtSignature, + }) +} + +// Constructors set the Type discriminant so callers never desync it. + +// NewTextContent returns a TextContent with the type discriminant set. +func NewTextContent(text string) TextContent { + return TextContent{Type: ContentTypeText, Text: text} +} + +// NewThinkingContent returns a ThinkingContent with the type discriminant set. +func NewThinkingContent(thinking string) ThinkingContent { + return ThinkingContent{Type: ContentTypeThinking, Thinking: thinking} +} + +// NewToolCallContent returns a ToolCallContent with the type discriminant set. +func NewToolCallContent(id, name string, arguments json.RawMessage) ToolCallContent { + return ToolCallContent{Type: ContentTypeToolCall, ID: id, Name: name, Arguments: arguments} +} + +// NewImageContent returns an ImageContent with the type discriminant set. +func NewImageContent(data, mimeType string) ImageContent { + return ImageContent{Type: ContentTypeImage, Data: data, MimeType: mimeType} +} + +// decodeContent peeks at the "type" field of a JSON object and decodes it into +// the matching concrete Content struct. This is the single dispatch point used +// by every container that holds Content (messages, tool results, session +// entries, provider parsing). +func decodeContent(raw json.RawMessage) (Content, error) { + var probe struct { + Type string `json:"type"` + } + if err := json.Unmarshal(raw, &probe); err != nil { + return nil, fmt.Errorf("content: peek type: %w", err) + } + switch probe.Type { + case ContentTypeText: + var c TextContent + if err := json.Unmarshal(raw, &c); err != nil { + return nil, err + } + return c, nil + case ContentTypeThinking: + var c ThinkingContent + if err := json.Unmarshal(raw, &c); err != nil { + return nil, err + } + return c, nil + case ContentTypeToolCall: + var c ToolCallContent + if err := json.Unmarshal(raw, &c); err != nil { + return nil, err + } + return c, nil + case ContentTypeImage: + var c ImageContent + if err := json.Unmarshal(raw, &c); err != nil { + return nil, err + } + return c, nil + case "": + return nil, fmt.Errorf("content: missing type discriminant") + default: + return nil, fmt.Errorf("content: unknown type %q", probe.Type) + } +} + +// ContentList is a slice of Content with discriminated JSON (un)marshalling. +// Fields typed []Content in messages use this so decoding dispatches on "type". +type ContentList []Content + +// UnmarshalJSON decodes a JSON array of content blocks, dispatching each element +// on its "type" discriminant. +func (cl *ContentList) UnmarshalJSON(data []byte) error { + var raws []json.RawMessage + if err := json.Unmarshal(data, &raws); err != nil { + return err + } + out := make(ContentList, 0, len(raws)) + for i, raw := range raws { + c, err := decodeContent(raw) + if err != nil { + return fmt.Errorf("content[%d]: %w", i, err) + } + out = append(out, c) + } + *cl = out + return nil +} diff --git a/pigo/internal/agentcore/event.go b/pigo/internal/agentcore/event.go new file mode 100644 index 0000000..1ad1126 --- /dev/null +++ b/pigo/internal/agentcore/event.go @@ -0,0 +1,217 @@ +package agentcore + +// AgentEvent is the sealed interface implemented by every event the loop emits. +// Consumers dispatch with a type switch, consistent with Content. pigo covers +// all 10 of pi's event types (PRD FR-24). +type AgentEvent interface { + isAgentEvent() + // EventType returns the discriminant string, useful for logging and for + // serialising events to the stream-json/stdio protocol (US-020). + EventType() string +} + +// Event type discriminants. +const ( + EventAgentStart = "agent_start" + EventAgentEnd = "agent_end" + EventTurnStart = "turn_start" + EventTurnEnd = "turn_end" + EventMessageStart = "message_start" + EventMessageUpdate = "message_update" + EventMessageEnd = "message_end" + EventToolExecutionStart = "tool_execution_start" + EventToolExecutionUpdate = "tool_execution_update" + EventToolExecutionEnd = "tool_execution_end" + EventCompaction = "compaction" + EventCompactionStart = "compaction_start" + EventTelemetry = "telemetry" + EventSubAgentProgress = "subagent_progress" +) + +// AgentStartEvent is emitted once when a loop run begins. SessionID, when set, +// is the id of the session backing this run; it is carried in the first +// stream-json event so a caller can associate output with a session and resume +// it later (mirrors pi/Claude Code, which put a session id in the first event). +type AgentStartEvent struct { + SessionID string +} + +// AgentEndEvent is emitted once when a loop run ends, carrying the messages +// newly produced during this run (the EventStream result). +type AgentEndEvent struct { + Messages []AgentMessage +} + +// TurnStartEvent marks the start of a turn (a single assistant response cycle). +type TurnStartEvent struct{} + +// TurnEndEvent marks the end of a turn, with the assistant message and any tool +// results produced during it. +type TurnEndEvent struct { + Message AssistantMessage + ToolResults []ToolResultMessage +} + +// MessageStartEvent is emitted when a message begins streaming. +type MessageStartEvent struct { + Message AgentMessage +} + +// MessageUpdateEvent is emitted for each streaming delta, carrying the current +// partial message and the raw provider-level event that produced it. +type MessageUpdateEvent struct { + Message AgentMessage + AssistantMessageEvent any +} + +// MessageEndEvent is emitted when a message finishes streaming. +type MessageEndEvent struct { + Message AgentMessage +} + +// ToolExecutionStartEvent is emitted before a tool runs. +type ToolExecutionStartEvent struct { + ToolCallID string + ToolName string + Args any +} + +// ToolExecutionUpdateEvent carries a partial result during tool execution. +type ToolExecutionUpdateEvent struct { + ToolCallID string + ToolName string + PartialResult AgentToolResult +} + +// ToolExecutionEndEvent is emitted when a tool finishes. +type ToolExecutionEndEvent struct { + ToolCallID string + ToolName string + Result AgentToolResult + IsError bool +} + +// CompactionEvent is emitted when the loop compacts the context window, either +// automatically (threshold/overflow) or on an explicit /compact request. It +// carries before/after token counts and how many messages were summarized vs. +// retained. When compaction fails it is still emitted with ErrorMessage set and +// the token/count fields describing the unchanged context, so consumers can +// surface the failure without the session aborting (US-004). +type CompactionEvent struct { + // Reason is why compaction ran: "manual", "threshold", or "overflow". + Reason string + // TokensBefore is the estimated context tokens prior to compaction. + TokensBefore int + // TokensAfter is the estimated context tokens after compaction (equals + // TokensBefore when compaction failed or was a no-op). + TokensAfter int + // SummarizedCount is the number of messages folded into the summary. + SummarizedCount int + // KeptCount is the number of recent messages retained verbatim. + KeptCount int + // ErrorMessage is non-empty when compaction failed; the original context is + // preserved in that case. + ErrorMessage string +} + +// CompactionStartEvent is emitted immediately before the loop runs compaction, +// so a front-end can show an in-progress "Compacting conversation…" indicator +// while the summarization request is in flight. The matching CompactionEvent is +// emitted when it completes (or fails). Reason mirrors CompactionEvent.Reason. +type CompactionStartEvent struct { + // Reason is why compaction is running: "manual", "threshold", or "overflow". + Reason string + // TokensBefore is the estimated context tokens that triggered compaction. + TokensBefore int +} + +// SubAgentProgressEvent carries structured progress from a running sub-agent +// (dispatched by the task tool). It is reported at the sub-agent's tool +// execution / turn boundaries so a TUI (multi-line status panel) or headless +// mode (stderr line) can display live progress. Elapsed time is intentionally +// omitted: consumers compute it themselves (TUI from tool-start time, headless +// from when the id was first seen) to avoid emitting an event per frame. +type SubAgentProgressEvent struct { + // ToolCallID is the parent task call's tool-call id, used as the key for + // the status line. + ToolCallID string + // Description is the task call's description, for display (may be empty). + Description string + // Activity is the current activity: tool name / phase, e.g. "Editing", + // "Running bash", "Thinking". + Activity string + // Tokens is the estimated sub-agent output token count (0 = unknown). + Tokens int +} + +// ToolTiming records how long one tool invocation took, keyed by tool name in +// TelemetryEvent.ToolDurationsMs. It aggregates repeated calls of the same tool +// so a summary stays compact regardless of turn count. +type ToolTiming struct { + // Count is how many times the tool was invoked over the run. + Count int + // TotalMs is the summed wall-clock duration of every invocation, in + // milliseconds. + TotalMs int64 +} + +// TelemetryEvent is a lightweight, additive observability summary emitted once +// at run end (just before agent_end) so scripts consuming the stream-json +// output can read structured metrics without a new dependency (no +// Prometheus/OTLP). It is purely observational: consumers that ignore it behave +// exactly as before. Metrics covered (observability — structured telemetry collection): +// - per-tool wall-clock durations (ToolDurationsMs, aggregated by tool name), +// - how many turns ran (Turns), +// - how many assistant responses were truncated by the output cap +// (TruncationCount), +// - how many times the context was compacted (CompactionCount), +// - the latest context-utilization ratio (ContextUtilization = used tokens / +// ContextWindow) and the raw numbers behind it. +type TelemetryEvent struct { + // Turns is the number of turns (turn_start events) the run executed. + Turns int + // ToolDurationsMs maps a tool name to its aggregated timing over the run. + ToolDurationsMs map[string]ToolTiming + // TruncationCount is how many assistant responses stopped with reason + // "length" (truncated by the output token cap), each triggering a resend. + TruncationCount int + // CompactionCount is how many successful context compactions occurred. + CompactionCount int + // ContextUtilization is the latest used/window ratio in [0,1], or 0 when the + // context window is unknown. Computed as ContextTokens / ContextWindow. + ContextUtilization float64 + // ContextTokens is the most recently observed estimated context-token usage. + ContextTokens int + // ContextWindow is the model's total context-token budget (0 when unknown). + ContextWindow int +} + +func (AgentStartEvent) isAgentEvent() {} +func (AgentEndEvent) isAgentEvent() {} +func (TurnStartEvent) isAgentEvent() {} +func (TurnEndEvent) isAgentEvent() {} +func (MessageStartEvent) isAgentEvent() {} +func (MessageUpdateEvent) isAgentEvent() {} +func (MessageEndEvent) isAgentEvent() {} +func (ToolExecutionStartEvent) isAgentEvent() {} +func (ToolExecutionUpdateEvent) isAgentEvent() {} +func (ToolExecutionEndEvent) isAgentEvent() {} +func (CompactionEvent) isAgentEvent() {} +func (CompactionStartEvent) isAgentEvent() {} +func (TelemetryEvent) isAgentEvent() {} +func (SubAgentProgressEvent) isAgentEvent() {} + +func (AgentStartEvent) EventType() string { return EventAgentStart } +func (AgentEndEvent) EventType() string { return EventAgentEnd } +func (TurnStartEvent) EventType() string { return EventTurnStart } +func (TurnEndEvent) EventType() string { return EventTurnEnd } +func (MessageStartEvent) EventType() string { return EventMessageStart } +func (MessageUpdateEvent) EventType() string { return EventMessageUpdate } +func (MessageEndEvent) EventType() string { return EventMessageEnd } +func (ToolExecutionStartEvent) EventType() string { return EventToolExecutionStart } +func (ToolExecutionUpdateEvent) EventType() string { return EventToolExecutionUpdate } +func (ToolExecutionEndEvent) EventType() string { return EventToolExecutionEnd } +func (CompactionEvent) EventType() string { return EventCompaction } +func (CompactionStartEvent) EventType() string { return EventCompactionStart } +func (TelemetryEvent) EventType() string { return EventTelemetry } +func (SubAgentProgressEvent) EventType() string { return EventSubAgentProgress } diff --git a/pigo/internal/agentcore/event_stream.go b/pigo/internal/agentcore/event_stream.go new file mode 100644 index 0000000..5a73341 --- /dev/null +++ b/pigo/internal/agentcore/event_stream.go @@ -0,0 +1,121 @@ +package agentcore + +import ( + "context" + "errors" + "sync" +) + +// EventStream is the Go equivalent of pi's EventStream: a producer pushes +// events onto a channel while a consumer ranges over them, and a terminal event +// yields a final result R. It replaces pi's async generator (event-stream.ts). +// +// Design (research §2.2): +// - Iteration: Events() returns <-chan T for `for ev := range s.Events()`. +// - Result: Result(ctx) blocks until the producer sets a result (or the +// stream fails/cancels). The result is NOT sent on the event channel, so a +// consumer that stops reading events can still obtain it. +// - Cancellation: the producer selects on ctx.Done() when sending, so a +// consumer that stops reading never leaks the producer goroutine. +// +// pi's isComplete/extractResult callbacks are retained as optional fields so a +// producer can let the stream detect the terminal event itself; a producer may +// instead call SetResult explicitly (more Go-idiomatic). Either path resolves +// Result exactly once. +type EventStream[T any, R any] struct { + ch chan T + + // IsComplete reports whether an event is the terminal one. Optional: if + // set, Emit auto-captures the result via ExtractResult when it returns true. + IsComplete func(event T) bool + // ExtractResult derives the final result from the terminal event. Required + // when IsComplete is set. + ExtractResult func(event T) R + + resultOnce sync.Once + result R + resultErr error + resultCh chan struct{} // closed once result (or resultErr) is set +} + +// ErrStreamIncomplete is returned by Result when the event channel closed +// without any result being set (the producer ended abnormally without a +// terminal event). +var ErrStreamIncomplete = errors.New("agent: event stream ended without a result") + +// NewEventStream constructs an EventStream with the given channel buffer size. +// A buffer of 0 gives fully synchronous back-pressure (each Emit blocks until a +// consumer receives), matching pi's sequential `await emit(...)`. +func NewEventStream[T any, R any](buffer int) *EventStream[T, R] { + if buffer < 0 { + buffer = 0 + } + return &EventStream[T, R]{ + ch: make(chan T, buffer), + resultCh: make(chan struct{}), + } +} + +// Events returns the receive-only event channel. Ranging over it terminates +// when the producer calls Close. +func (s *EventStream[T, R]) Events() <-chan T { return s.ch } + +// Emit sends an event to consumers, honoring cancellation. If ctx is cancelled +// before the event is received, Emit returns ctx.Err() and the event is +// dropped. When IsComplete is configured and reports true for the event, the +// result is captured (once) before the send. +func (s *EventStream[T, R]) Emit(ctx context.Context, event T) error { + if s.IsComplete != nil && s.IsComplete(event) && s.ExtractResult != nil { + s.SetResult(s.ExtractResult(event)) + } + select { + case s.ch <- event: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// SetResult records the final result. Only the first call wins; later calls +// (including SetError) are no-ops. Safe to call before Close. +func (s *EventStream[T, R]) SetResult(result R) { + s.resultOnce.Do(func() { + s.result = result + close(s.resultCh) + }) +} + +// SetError records a terminal error as the stream's outcome. Only the first +// call among SetResult/SetError wins. +func (s *EventStream[T, R]) SetError(err error) { + s.resultOnce.Do(func() { + s.resultErr = err + close(s.resultCh) + }) +} + +// Close closes the event channel, ending consumer iteration. If no result was +// set, Result will report ErrStreamIncomplete. Call exactly once from the +// producer after the last Emit. +func (s *EventStream[T, R]) Close() { + // Ensure a waiting Result never blocks forever if the producer forgot to + // set a result. + s.resultOnce.Do(func() { + s.resultErr = ErrStreamIncomplete + close(s.resultCh) + }) + close(s.ch) +} + +// Result blocks until the producer sets a result/error, ctx is cancelled, or +// the stream closes without a result. It is safe to call concurrently and +// returns the same outcome on every call. +func (s *EventStream[T, R]) Result(ctx context.Context) (R, error) { + select { + case <-s.resultCh: + return s.result, s.resultErr + case <-ctx.Done(): + var zero R + return zero, ctx.Err() + } +} diff --git a/pigo/internal/agentcore/event_stream_test.go b/pigo/internal/agentcore/event_stream_test.go new file mode 100644 index 0000000..d12d2b5 --- /dev/null +++ b/pigo/internal/agentcore/event_stream_test.go @@ -0,0 +1,121 @@ +package agentcore + +import ( + "context" + "testing" + "time" +) + +// TestEventStreamNormalCompletion drives a producer that emits events and sets +// a result, then verifies the consumer sees every event and Result yields the +// captured value. +func TestEventStreamNormalCompletion(t *testing.T) { + s := NewEventStream[AgentEvent, []AgentMessage](0) + want := []AgentMessage{ + UserMessage{RoleField: RoleUser, Content: ContentList{NewTextContent("hi")}}, + } + + go func() { + ctx := context.Background() + _ = s.Emit(ctx, TurnStartEvent{}) + _ = s.Emit(ctx, AgentEndEvent{Messages: want}) + s.SetResult(want) + s.Close() + }() + + var got int + for range s.Events() { + got++ + } + if got != 2 { + t.Fatalf("want 2 events, got %d", got) + } + res, err := s.Result(context.Background()) + if err != nil { + t.Fatalf("result: %v", err) + } + if len(res) != 1 || res[0].Role() != RoleUser { + t.Fatalf("result payload wrong: %+v", res) + } +} + +// TestEventStreamIsCompleteCallback verifies the isComplete/extractResult +// callbacks auto-capture the result on the terminal event. +func TestEventStreamIsCompleteCallback(t *testing.T) { + s := NewEventStream[AgentEvent, []AgentMessage](4) + s.IsComplete = func(e AgentEvent) bool { return e.EventType() == EventAgentEnd } + s.ExtractResult = func(e AgentEvent) []AgentMessage { return e.(AgentEndEvent).Messages } + + msgs := []AgentMessage{AssistantMessage{RoleField: RoleAssistant}} + go func() { + ctx := context.Background() + _ = s.Emit(ctx, MessageStartEvent{}) + _ = s.Emit(ctx, AgentEndEvent{Messages: msgs}) + s.Close() + }() + + for range s.Events() { + } + res, err := s.Result(context.Background()) + if err != nil { + t.Fatalf("result: %v", err) + } + if len(res) != 1 { + t.Fatalf("want 1 msg from extractResult, got %d", len(res)) + } +} + +// TestEventStreamCancellation verifies that a cancelled context unblocks a +// producer stuck on Emit (consumer stopped reading) and that Result returns the +// context error. +func TestEventStreamCancellation(t *testing.T) { + s := NewEventStream[AgentEvent, []AgentMessage](0) + ctx, cancel := context.WithCancel(context.Background()) + + emitErr := make(chan error, 1) + go func() { + // First emit has no consumer; it blocks until cancel. + emitErr <- s.Emit(ctx, TurnStartEvent{}) + }() + + // Give the producer a moment to block on the send, then cancel. + time.Sleep(20 * time.Millisecond) + cancel() + + select { + case err := <-emitErr: + if err == nil { + t.Fatal("expected Emit to return ctx error on cancellation") + } + case <-time.After(time.Second): + t.Fatal("Emit did not unblock after cancel (goroutine leak)") + } + + // Result with a cancelled context returns promptly with the ctx error. + if _, err := s.Result(ctx); err == nil { + t.Fatal("expected Result to return ctx error") + } +} + +// TestEventStreamIncompleteClose verifies Close without a result yields +// ErrStreamIncomplete. +func TestEventStreamIncompleteClose(t *testing.T) { + s := NewEventStream[AgentEvent, []AgentMessage](1) + s.Close() + if _, err := s.Result(context.Background()); err != ErrStreamIncomplete { + t.Fatalf("want ErrStreamIncomplete, got %v", err) + } +} + +// TestEventStreamSetErrorWins verifies SetError is reported and later SetResult +// is ignored (first outcome wins). +func TestEventStreamSetErrorWins(t *testing.T) { + s := NewEventStream[AgentEvent, []AgentMessage](1) + sentinel := context.Canceled + s.SetError(sentinel) + s.SetResult(nil) + s.Close() + if _, err := s.Result(context.Background()); err != sentinel { + t.Fatalf("want sentinel error, got %v", err) + } +} diff --git a/pigo/internal/agentcore/helpers.go b/pigo/internal/agentcore/helpers.go new file mode 100644 index 0000000..d600bd3 --- /dev/null +++ b/pigo/internal/agentcore/helpers.go @@ -0,0 +1,60 @@ +package agentcore + +import ( + "context" + "encoding/json" + "strings" +) + +// ContentToText flattens text blocks of a content list into a single string, +// the lowest-common-denominator representation accepted by every OpenAI- +// compatible gateway. Non-text blocks (thinking, tool calls) are surfaced +// through their own fields, so they are skipped here. +func ContentToText(list ContentList) string { + var b strings.Builder + for _, c := range list { + if tc, ok := c.(TextContent); ok { + b.WriteString(tc.Text) + } + } + return b.String() +} + +// LastAssistantOf returns a pointer to the last AssistantMessage in msgs, or nil. +func LastAssistantOf(msgs []AgentMessage) *AssistantMessage { + for i := len(msgs) - 1; i >= 0; i-- { + if a, ok := msgs[i].(AssistantMessage); ok { + return &a + } + } + return nil +} + +// EmitFunc emits a loop-level AgentEvent honoring cancellation. +type EmitFunc func(ctx context.Context, ev AgentEvent) error + +// PrepareArgumentsFunc optionally rewrites a tool's raw arguments before schema +// validation (e.g. injecting defaults). An error aborts the call with an error +// result. Optional (nil = identity). +type PrepareArgumentsFunc func(ctx context.Context, toolName string, args json.RawMessage) (json.RawMessage, error) + +// BeforeToolCallDecision is the optional result of the beforeToolCall hook. When +// Block is true the tool is not executed and an error result is produced; +// Content/Details override the default block message when set. When Block is +// false and UpdatedInput is non-empty, it replaces the tool's raw arguments +// before execution (PreToolUse rewrite, FR-8); the replacement is re-validated +// against the tool schema. +type BeforeToolCallDecision struct { + Block bool + Content *ContentList + Details *any + UpdatedInput json.RawMessage +} + +// BeforeToolCallFunc runs after validation and may block the call (permission / +// sandbox checks, FR-4/FR-26). Returning nil allows the call. Optional. +type BeforeToolCallFunc func(ctx context.Context, call AgentToolCall) *BeforeToolCallDecision + +// AfterToolCallFunc runs after execution and may override the result +// field-by-field via AfterToolCallResult (FR-5, no deep merge). Optional. +type AfterToolCallFunc func(ctx context.Context, call AgentToolCall, result AgentToolResult, isError bool) *AfterToolCallResult diff --git a/pigo/internal/agentcore/helpers_test.go b/pigo/internal/agentcore/helpers_test.go new file mode 100644 index 0000000..320a0ef --- /dev/null +++ b/pigo/internal/agentcore/helpers_test.go @@ -0,0 +1,141 @@ +package agentcore + +import ( + "encoding/json" + "testing" +) + +func TestContentToText(t *testing.T) { + cases := []struct { + name string + list ContentList + want string + }{ + {"empty", nil, ""}, + {"single text", ContentList{NewTextContent("hello")}, "hello"}, + { + "skips non-text blocks", + ContentList{ + NewTextContent("a"), + NewThinkingContent("ignored"), + NewToolCallContent("c1", "ls", json.RawMessage(`{}`)), + NewTextContent("b"), + NewImageContent("data", "image/png"), + }, + "ab", + }, + { + "only non-text", + ContentList{NewThinkingContent("x"), NewImageContent("d", "image/png")}, + "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := ContentToText(tc.list); got != tc.want { + t.Errorf("ContentToText = %q, want %q", got, tc.want) + } + }) + } +} + +func TestLastAssistantOf(t *testing.T) { + t.Run("nil when absent", func(t *testing.T) { + msgs := []AgentMessage{ + UserMessage{RoleField: RoleUser, Content: ContentList{NewTextContent("hi")}}, + ToolResultMessage{RoleField: RoleToolResult, ToolCallID: "c1"}, + } + if got := LastAssistantOf(msgs); got != nil { + t.Errorf("want nil, got %+v", got) + } + }) + t.Run("nil for empty slice", func(t *testing.T) { + if got := LastAssistantOf(nil); got != nil { + t.Errorf("want nil, got %+v", got) + } + }) + t.Run("returns last assistant", func(t *testing.T) { + msgs := []AgentMessage{ + AssistantMessage{RoleField: RoleAssistant, Model: "first"}, + UserMessage{RoleField: RoleUser}, + AssistantMessage{RoleField: RoleAssistant, Model: "last"}, + ToolResultMessage{RoleField: RoleToolResult}, + } + got := LastAssistantOf(msgs) + if got == nil { + t.Fatal("want an assistant message, got nil") + } + if got.Model != "last" { + t.Errorf("want the last assistant (model %q), got %q", "last", got.Model) + } + }) +} + +func TestToolCallsEmpty(t *testing.T) { + m := AssistantMessage{ + RoleField: RoleAssistant, + Content: ContentList{NewTextContent("no tools here"), NewThinkingContent("hmm")}, + } + if calls := m.ToolCalls(); calls != nil { + t.Errorf("want nil for a message with no tool calls, got %+v", calls) + } +} + +func TestToolCallsPreservesOrder(t *testing.T) { + m := AssistantMessage{ + RoleField: RoleAssistant, + Content: ContentList{ + NewToolCallContent("c1", "read", json.RawMessage(`{}`)), + NewTextContent("between"), + NewToolCallContent("c2", "write", json.RawMessage(`{}`)), + }, + } + calls := m.ToolCalls() + if len(calls) != 2 { + t.Fatalf("want 2 tool calls, got %d", len(calls)) + } + if calls[0].ID != "c1" || calls[1].ID != "c2" { + t.Errorf("tool call order lost: %+v", calls) + } +} + +// TestNewContentConstructorsSetType guards the invariant that every constructor +// sets its type discriminant, so a marshalled block always carries a "type". +func TestNewContentConstructorsSetType(t *testing.T) { + cases := []struct { + got Content + want string + }{ + {NewTextContent("t"), ContentTypeText}, + {NewThinkingContent("th"), ContentTypeThinking}, + {NewToolCallContent("id", "n", json.RawMessage(`{}`)), ContentTypeToolCall}, + {NewImageContent("d", "image/png"), ContentTypeImage}, + } + for _, tc := range cases { + data, err := json.Marshal(tc.got) + if err != nil { + t.Fatalf("marshal %T: %v", tc.got, err) + } + var probe struct { + Type string `json:"type"` + } + if err := json.Unmarshal(data, &probe); err != nil { + t.Fatalf("unmarshal probe %T: %v", tc.got, err) + } + if probe.Type != tc.want { + t.Errorf("%T type = %q, want %q", tc.got, probe.Type, tc.want) + } + } +} + +func TestRoleAccessors(t *testing.T) { + if got := (UserMessage{}).Role(); got != RoleUser { + t.Errorf("UserMessage.Role = %q, want %q", got, RoleUser) + } + if got := (AssistantMessage{}).Role(); got != RoleAssistant { + t.Errorf("AssistantMessage.Role = %q, want %q", got, RoleAssistant) + } + if got := (ToolResultMessage{}).Role(); got != RoleToolResult { + t.Errorf("ToolResultMessage.Role = %q, want %q", got, RoleToolResult) + } +} diff --git a/pigo/internal/agentcore/hooks.go b/pigo/internal/agentcore/hooks.go new file mode 100644 index 0000000..9f03dce --- /dev/null +++ b/pigo/internal/agentcore/hooks.go @@ -0,0 +1,44 @@ +package agentcore + +// ThinkingLevel is the unified reasoning-effort enum (agent layer). It keeps +// pi's full 6 levels; providers map it to their own wire format via a +// per-model ThinkingLevelMap (decision #10). +type ThinkingLevel string + +const ( + ThinkingOff ThinkingLevel = "off" + ThinkingMinimal ThinkingLevel = "minimal" + ThinkingLow ThinkingLevel = "low" + ThinkingMedium ThinkingLevel = "medium" + ThinkingHigh ThinkingLevel = "high" + ThinkingXHigh ThinkingLevel = "xhigh" + ThinkingMax ThinkingLevel = "max" +) + +// ThinkingLevelMap maps a unified level to a provider-specific wire value. +// A nil value means "supported but disabled at this level"; an absent key means +// "this level is not supported by the model". The pointer is what distinguishes +// those two cases, so it must stay *string. +type ThinkingLevelMap map[ThinkingLevel]*string + +// AfterToolCallResult is the optional override returned by the afterToolCall +// hook. Every field is a pointer so the loop can distinguish "not provided" +// (nil) from "provided, possibly zero" — pi expresses this with `??`, Go needs +// pointers. Fields are applied with field-level replacement, no deep merge +// (FR-5). +type AfterToolCallResult struct { + Content *ContentList + Details *any + Terminate *bool + IsError *bool +} + +// AgentLoopTurnUpdate is the optional result of the prepareNextTurn hook: it can +// swap the context, model, or thinking level for the next turn. Pointer fields +// distinguish "not provided" from an explicit value; ThinkingLevel is +// three-state (nil = keep, &"off" = disable, &level = set). +type AgentLoopTurnUpdate struct { + Context *AgentContext + Model *string + ThinkingLevel *ThinkingLevel +} diff --git a/pigo/internal/agentcore/message.go b/pigo/internal/agentcore/message.go new file mode 100644 index 0000000..4204843 --- /dev/null +++ b/pigo/internal/agentcore/message.go @@ -0,0 +1,210 @@ +package agentcore + +import ( + "encoding/json" + "fmt" +) + +// Message roles, matching pi's wire format. +const ( + RoleUser = "user" + RoleAssistant = "assistant" + RoleToolResult = "toolResult" + // RoleCompaction marks a compaction checkpoint persisted inline in the + // message list: it replaces the history summarized before it (pi's + // "compactionSummary"). It is not sent to the model verbatim; the LLM + // conversion turns it into a user text block. + RoleCompaction = "compaction" +) + +// Message is the sealed interface implemented by the three message roles. +// AgentMessage (the loop's message abstraction) is simply Message: custom +// message kinds implement the same interface and convertToLlm filters out any +// that are not LLM-bound. This deliberately replaces pi's declaration merging, +// which has no Go equivalent. +type Message interface { + isMessage() + // Role returns the discriminant ("user" | "assistant" | "toolResult"). + Role() string +} + +// AgentMessage is the loop-level message type. It is the same as Message; the +// alias documents intent at call sites that deal with the loop rather than raw +// LLM messages. +type AgentMessage = Message + +// Usage reports token accounting for an assistant response. +type Usage struct { + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` +} + +// UserMessage is input from the user. Content is restricted at construction to +// text/image blocks (runtime constraint, not a separate interface). +type UserMessage struct { + RoleField string `json:"role"` + Content ContentList `json:"content"` + Timestamp int64 `json:"timestamp"` +} + +func (UserMessage) isMessage() {} +func (m UserMessage) Role() string { return RoleUser } + +// AssistantMessage is a model response. Content may hold text/thinking/toolCall +// blocks. StopReason follows pi's set (end_turn/tool_use/length/error/aborted). +type AssistantMessage struct { + RoleField string `json:"role"` + Content ContentList `json:"content"` + API string `json:"api,omitempty"` + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` + Usage *Usage `json:"usage,omitempty"` + StopReason string `json:"stopReason,omitempty"` + ErrorMessage string `json:"errorMessage,omitempty"` + Timestamp int64 `json:"timestamp"` + + // Optional diagnostics, kept for cross-provider replay/observability. + ResponseModel string `json:"responseModel,omitempty"` + ResponseID string `json:"responseId,omitempty"` +} + +func (AssistantMessage) isMessage() {} +func (m AssistantMessage) Role() string { return RoleAssistant } + +// ToolCalls returns the tool call blocks in this assistant message, in order. +func (m AssistantMessage) ToolCalls() []ToolCallContent { + var calls []ToolCallContent + for _, c := range m.Content { + if tc, ok := c.(ToolCallContent); ok { + calls = append(calls, tc) + } + } + return calls +} + +// ToolResultMessage carries the outcome of executing a single tool call. +// Content is restricted to text/image blocks at construction. +type ToolResultMessage struct { + RoleField string `json:"role"` + ToolCallID string `json:"toolCallId"` + ToolName string `json:"toolName"` + Content ContentList `json:"content"` + Details any `json:"details,omitempty"` + IsError bool `json:"isError"` + Timestamp int64 `json:"timestamp"` +} + +func (ToolResultMessage) isMessage() {} +func (m ToolResultMessage) Role() string { return RoleToolResult } + +// CompactionMessage is a summarization checkpoint persisted inline in the +// message list. It stands in for the history compacted before it: Summary is +// the structured checkpoint text and TokensBefore records the estimated context +// size at compaction time (for observability). Details optionally holds the +// file operations extracted from the compacted range. Mirrors pi's +// CompactionSummaryMessage + CompactionEntry. +type CompactionMessage struct { + RoleField string `json:"role"` + Summary string `json:"summary"` + TokensBefore int `json:"tokensBefore,omitempty"` + // Details is opaque at this layer (the compaction package owns its shape); + // kept as raw JSON so agentcore stays free of a compaction dependency. + Details json.RawMessage `json:"details,omitempty"` + Timestamp int64 `json:"timestamp"` +} + +func (CompactionMessage) isMessage() {} +func (m CompactionMessage) Role() string { return RoleCompaction } + +// compactionSummaryPrefix / compactionSummarySuffix wrap a compaction summary +// when it is rendered into an LLM user message, matching pi's +// COMPACTION_SUMMARY_PREFIX / COMPACTION_SUMMARY_SUFFIX. +const ( + compactionSummaryPrefix = "The conversation history before this point was compacted into the following summary:\n\n\n" + compactionSummarySuffix = "\n" +) + +// AsUserMessage renders a compaction checkpoint as the user text message that +// stands in for the compacted history when building the LLM request. The +// provider encoders call this so a persisted compaction line replays as +// context rather than being dropped. +func (m CompactionMessage) AsUserMessage() UserMessage { + return UserMessage{ + RoleField: RoleUser, + Content: ContentList{NewTextContent(compactionSummaryPrefix + m.Summary + compactionSummarySuffix)}, + Timestamp: m.Timestamp, + } +} + +// StopReason values, matching pi. +const ( + StopReasonEndTurn = "end_turn" + StopReasonToolUse = "tool_use" + StopReasonLength = "length" + StopReasonError = "error" + StopReasonAborted = "aborted" +) + +// MessageList is a slice of Message with discriminated JSON (un)marshalling, +// dispatching on the "role" field. Used by AgentContext and session persistence. +type MessageList []Message + +// UnmarshalJSON decodes a JSON array of messages, dispatching each element on +// its "role" discriminant. +func (ml *MessageList) UnmarshalJSON(data []byte) error { + var raws []json.RawMessage + if err := json.Unmarshal(data, &raws); err != nil { + return err + } + out := make(MessageList, 0, len(raws)) + for i, raw := range raws { + m, err := decodeMessage(raw) + if err != nil { + return fmt.Errorf("message[%d]: %w", i, err) + } + out = append(out, m) + } + *ml = out + return nil +} + +// decodeMessage peeks at the "role" field and decodes into the matching +// concrete message struct. +func decodeMessage(raw json.RawMessage) (Message, error) { + var probe struct { + Role string `json:"role"` + } + if err := json.Unmarshal(raw, &probe); err != nil { + return nil, fmt.Errorf("peek role: %w", err) + } + switch probe.Role { + case RoleUser: + var m UserMessage + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil + case RoleAssistant: + var m AssistantMessage + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil + case RoleToolResult: + var m ToolResultMessage + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil + case RoleCompaction: + var m CompactionMessage + if err := json.Unmarshal(raw, &m); err != nil { + return nil, err + } + return m, nil + case "": + return nil, fmt.Errorf("missing role discriminant") + default: + return nil, fmt.Errorf("unknown role %q", probe.Role) + } +} diff --git a/pigo/internal/agentcore/progress_ctx.go b/pigo/internal/agentcore/progress_ctx.go new file mode 100644 index 0000000..b223bf6 --- /dev/null +++ b/pigo/internal/agentcore/progress_ctx.go @@ -0,0 +1,22 @@ +package agentcore + +import "context" + +// progressEmitterKey is the unexported context key under which a run-level +// progress EmitFunc is stored. +type progressEmitterKey struct{} + +// WithProgressEmitter returns a child context carrying emit as the run-level +// progress emitter. The task tool injects the parent loop's EmitFunc here so a +// dispatched sub-agent can surface SubAgentProgressEvent up the parent stream. +func WithProgressEmitter(ctx context.Context, emit EmitFunc) context.Context { + return context.WithValue(ctx, progressEmitterKey{}, emit) +} + +// ProgressEmitterFromContext returns the run-level progress emitter carried by +// ctx, or nil if none was set (in which case callers should skip progress +// reporting rather than panic). +func ProgressEmitterFromContext(ctx context.Context) EmitFunc { + emit, _ := ctx.Value(progressEmitterKey{}).(EmitFunc) + return emit +} diff --git a/pigo/internal/agentcore/progress_ctx_test.go b/pigo/internal/agentcore/progress_ctx_test.go new file mode 100644 index 0000000..9c64574 --- /dev/null +++ b/pigo/internal/agentcore/progress_ctx_test.go @@ -0,0 +1,51 @@ +package agentcore + +import ( + "context" + "errors" + "testing" +) + +func TestSubAgentProgressEventImplementsAgentEvent(t *testing.T) { + var ev AgentEvent = SubAgentProgressEvent{ + ToolCallID: "call-1", + Description: "do a thing", + Activity: "Editing", + Tokens: 42, + } + if got := ev.EventType(); got != EventSubAgentProgress { + t.Fatalf("EventType() = %q, want %q", got, EventSubAgentProgress) + } + if EventSubAgentProgress != "subagent_progress" { + t.Fatalf("EventSubAgentProgress = %q, want %q", EventSubAgentProgress, "subagent_progress") + } +} + +func TestProgressEmitterRoundTrip(t *testing.T) { + var seen AgentEvent + sentinel := errors.New("sentinel") + emit := func(ctx context.Context, ev AgentEvent) error { + seen = ev + return sentinel + } + + ctx := WithProgressEmitter(context.Background(), emit) + got := ProgressEmitterFromContext(ctx) + if got == nil { + t.Fatal("ProgressEmitterFromContext returned nil after WithProgressEmitter") + } + + want := SubAgentProgressEvent{ToolCallID: "call-2", Activity: "Thinking"} + if err := got(ctx, want); !errors.Is(err, sentinel) { + t.Fatalf("emitter returned err = %v, want sentinel", err) + } + if seen != want { + t.Fatalf("emitter received %#v, want %#v", seen, want) + } +} + +func TestProgressEmitterFromBareContextIsNil(t *testing.T) { + if got := ProgressEmitterFromContext(context.Background()); got != nil { + t.Fatalf("ProgressEmitterFromContext on bare ctx = %v, want nil", got) + } +} diff --git a/pigo/internal/agentcore/tool.go b/pigo/internal/agentcore/tool.go new file mode 100644 index 0000000..5fc5d86 --- /dev/null +++ b/pigo/internal/agentcore/tool.go @@ -0,0 +1,61 @@ +package agentcore + +import ( + "context" + "encoding/json" +) + +// AgentContext is the input state for a loop run: system prompt, conversation +// messages, and the tools available to the model. +type AgentContext struct { + SystemPrompt string `json:"systemPrompt"` + Messages MessageList `json:"messages"` + Tools []AgentTool `json:"-"` +} + +// ToolExecutionMode selects how a tool is executed relative to others in a batch. +type ToolExecutionMode string + +const ( + // ToolExecutionParallel allows the tool to run concurrently with others. + ToolExecutionParallel ToolExecutionMode = "parallel" + // ToolExecutionSequential forces the whole batch to run serially. + ToolExecutionSequential ToolExecutionMode = "sequential" +) + +// ToolUpdateFunc receives a partial result during tool execution; the loop +// turns each call into a tool_execution_update event. +type ToolUpdateFunc func(partial AgentToolResult) + +// AgentTool is a tool the model can invoke. Schema is the JSON Schema used to +// validate arguments before execution (US-014). +type AgentTool interface { + Name() string + Description() string + // Schema returns the JSON Schema (as raw JSON) for the tool's arguments. + Schema() json.RawMessage + // ExecutionMode reports whether this tool forces sequential execution. + ExecutionMode() ToolExecutionMode + // Execute runs the tool. onUpdate may be nil. + Execute(ctx context.Context, id string, args json.RawMessage, onUpdate ToolUpdateFunc) (AgentToolResult, error) +} + +// AgentToolCall is a decoded request to invoke a tool (the loop-level view of a +// ToolCallContent block). +type AgentToolCall struct { + ID string `json:"id"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` +} + +// AgentToolResult is the outcome of executing a tool. +// +// Details uses `any` in the first version (matching pi's internal +// AgentToolResult); a generic form can be added later. Terminate is a +// *bool so "not set" is distinguishable from an explicit false — the loop only +// signals early termination when every result in a batch has Terminate=true. +type AgentToolResult struct { + Content ContentList `json:"content"` + Details any `json:"details,omitempty"` + Terminate *bool `json:"terminate,omitempty"` +} diff --git a/pigo/internal/agentcore/types_test.go b/pigo/internal/agentcore/types_test.go new file mode 100644 index 0000000..6afb401 --- /dev/null +++ b/pigo/internal/agentcore/types_test.go @@ -0,0 +1,182 @@ +package agentcore + +import ( + "encoding/json" + "testing" +) + +func TestContentListRoundTrip(t *testing.T) { + in := ContentList{ + NewTextContent("hello"), + NewThinkingContent("pondering"), + NewToolCallContent("call_1", "read", json.RawMessage(`{"path":"a.go"}`)), + NewImageContent("YmFzZTY0", "image/png"), + } + data, err := json.Marshal(in) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var out ContentList + if err := json.Unmarshal(data, &out); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if len(out) != 4 { + t.Fatalf("want 4 blocks, got %d", len(out)) + } + if _, ok := out[0].(TextContent); !ok { + t.Errorf("block 0: want TextContent, got %T", out[0]) + } + if _, ok := out[1].(ThinkingContent); !ok { + t.Errorf("block 1: want ThinkingContent, got %T", out[1]) + } + tc, ok := out[2].(ToolCallContent) + if !ok { + t.Fatalf("block 2: want ToolCallContent, got %T", out[2]) + } + if tc.ID != "call_1" || tc.Name != "read" { + t.Errorf("toolCall fields lost: %+v", tc) + } + if string(tc.Arguments) != `{"path":"a.go"}` { + t.Errorf("arguments lost: %s", tc.Arguments) + } + if _, ok := out[3].(ImageContent); !ok { + t.Errorf("block 3: want ImageContent, got %T", out[3]) + } +} + +func TestContentUnknownTypeRejected(t *testing.T) { + var out ContentList + err := json.Unmarshal([]byte(`[{"type":"bogus"}]`), &out) + if err == nil { + t.Fatal("expected error for unknown content type") + } +} + +// TestToolCallInvalidArgumentsMarshal verifies a ToolCallContent whose +// Arguments are syntactically invalid JSON (as a model can stream) still +// marshals — as a JSON string of the raw bytes — rather than aborting the +// encode. Without this, session persistence and provider re-serialization would +// crash the whole turn on a single malformed tool call. +func TestToolCallInvalidArgumentsMarshal(t *testing.T) { + bad := NewToolCallContent("c1", "todo", json.RawMessage(`{"todos": []{}"content": ""x"}`)) + data, err := json.Marshal(bad) + if err != nil { + t.Fatalf("marshal invalid tool args: %v", err) + } + if !json.Valid(data) { + t.Fatalf("marshaled output is not valid JSON: %s", data) + } + // It must round-trip back through the discriminated decoder without error. + var out ContentList + if err := json.Unmarshal([]byte("["+string(data)+"]"), &out); err != nil { + t.Fatalf("round-trip unmarshal: %v", err) + } + tc, ok := out[0].(ToolCallContent) + if !ok { + t.Fatalf("want ToolCallContent, got %T", out[0]) + } + // The raw invalid text is preserved (as the decoded string). + var recovered string + if err := json.Unmarshal(tc.Arguments, &recovered); err != nil { + t.Fatalf("arguments not a JSON string: %v", err) + } + if recovered != `{"todos": []{}"content": ""x"}` { + t.Errorf("raw arguments lost: %q", recovered) + } +} + +// TestToolCallValidArgumentsUnchanged verifies well-formed arguments are emitted +// verbatim (not string-wrapped), preserving the object shape providers expect. +func TestToolCallValidArgumentsUnchanged(t *testing.T) { + tc := NewToolCallContent("c1", "read", json.RawMessage(`{"path":"a.go"}`)) + data, err := json.Marshal(tc) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var out ContentList + if err := json.Unmarshal([]byte("["+string(data)+"]"), &out); err != nil { + t.Fatalf("unmarshal: %v", err) + } + got := out[0].(ToolCallContent) + if string(got.Arguments) != `{"path":"a.go"}` { + t.Errorf("arguments = %s, want the object unchanged", got.Arguments) + } +} + +func TestContentMissingTypeRejected(t *testing.T) { + var out ContentList + err := json.Unmarshal([]byte(`[{"text":"no type"}]`), &out) + if err == nil { + t.Fatal("expected error for missing type discriminant") + } +} + +func TestMessageListRoundTrip(t *testing.T) { + term := true + _ = term + in := MessageList{ + UserMessage{RoleField: RoleUser, Content: ContentList{NewTextContent("hi")}, Timestamp: 1}, + AssistantMessage{ + RoleField: RoleAssistant, + Content: ContentList{NewTextContent("ok"), NewToolCallContent("c1", "ls", json.RawMessage(`{}`))}, + StopReason: StopReasonToolUse, + Timestamp: 2, + }, + ToolResultMessage{RoleField: RoleToolResult, ToolCallID: "c1", ToolName: "ls", Content: ContentList{NewTextContent("file.go")}, Timestamp: 3}, + } + data, err := json.Marshal(in) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var out MessageList + if err := json.Unmarshal(data, &out); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if len(out) != 3 { + t.Fatalf("want 3 messages, got %d", len(out)) + } + if out[0].Role() != RoleUser { + t.Errorf("msg 0: want user, got %s", out[0].Role()) + } + am, ok := out[1].(AssistantMessage) + if !ok { + t.Fatalf("msg 1: want AssistantMessage, got %T", out[1]) + } + if calls := am.ToolCalls(); len(calls) != 1 || calls[0].Name != "ls" { + t.Errorf("assistant ToolCalls wrong: %+v", calls) + } + if out[2].Role() != RoleToolResult { + t.Errorf("msg 2: want toolResult, got %s", out[2].Role()) + } +} + +func TestMessageUnknownRoleRejected(t *testing.T) { + var out MessageList + if err := json.Unmarshal([]byte(`[{"role":"system"}]`), &out); err == nil { + t.Fatal("expected error for unknown role") + } +} + +// TestAgentEventCoverage asserts all 10 event types report a distinct, +// non-empty discriminant (PRD FR-24). +func TestAgentEventCoverage(t *testing.T) { + events := []AgentEvent{ + AgentStartEvent{}, AgentEndEvent{}, TurnStartEvent{}, TurnEndEvent{}, + MessageStartEvent{}, MessageUpdateEvent{}, MessageEndEvent{}, + ToolExecutionStartEvent{}, ToolExecutionUpdateEvent{}, ToolExecutionEndEvent{}, + } + seen := map[string]bool{} + for _, e := range events { + et := e.EventType() + if et == "" { + t.Errorf("%T has empty EventType", e) + } + if seen[et] { + t.Errorf("duplicate event type %q", et) + } + seen[et] = true + } + if len(seen) != 10 { + t.Fatalf("want 10 distinct event types, got %d", len(seen)) + } +} diff --git a/pigo/internal/agenttool/bash_background.go b/pigo/internal/agenttool/bash_background.go new file mode 100644 index 0000000..ef9dfd2 --- /dev/null +++ b/pigo/internal/agenttool/bash_background.go @@ -0,0 +1,167 @@ +// This file implements background bash execution: a BashJobStore holding the +// commands launched with run_in_background, and the BashJob state each one +// carries. A background job is detached from the turn context (which is canceled +// when the turn ends) and runs under its own cancelable context until it exits or +// kill_bash stops it. Its combined stdout/stderr accumulates in a buffer that +// bash_output drains incrementally, mirroring Claude Code's background shells + +// BashOutput/KillShell. +package agenttool + +import ( + "bytes" + "context" + "fmt" + "sync" + "time" +) + +// BashJobStatus is the lifecycle state of a background bash job. +type BashJobStatus string + +const ( + // BashRunning means the command is still executing. + BashRunning BashJobStatus = "running" + // BashExited means the command finished (successfully or not) or was killed. + BashExited BashJobStatus = "exited" +) + +// BashJob is a single background command: its identity, growing combined output, +// and terminal status. All fields are guarded by mu so the running command's +// writer, bash_output reads, and kill_bash can touch it concurrently. +type BashJob struct { + // ID is the stable handle (e.g. "bash_1") bash_output/kill_bash address. + ID string + // Command is the shell command line, kept for listing/display. + Command string + // StartedAt is when the command was launched. + StartedAt time.Time + + mu sync.Mutex + buf bytes.Buffer + cursor int // bytes of buf already returned by bash_output + status BashJobStatus + exitCode int + errMsg string + finished time.Time + cancel context.CancelFunc +} + +// jobWriter adapts a BashJob to io.Writer so it can be a command's Stdout/Stderr; +// each write appends to the job's combined buffer under its lock. +type jobWriter struct{ job *BashJob } + +func (w jobWriter) Write(p []byte) (int, error) { + w.job.mu.Lock() + w.job.buf.Write(p) + w.job.mu.Unlock() + return len(p), nil +} + +// writer returns an io.Writer that appends to the job's output buffer. +func (j *BashJob) writer() jobWriter { return jobWriter{job: j} } + +// finish records the command's terminal state from the error returned by +// cmd.Wait (nil = success). It is idempotent-safe to call once per job. +func (j *BashJob) finish(exitCode int, errMsg string) { + j.mu.Lock() + defer j.mu.Unlock() + j.status = BashExited + j.exitCode = exitCode + j.errMsg = errMsg + j.finished = time.Now() +} + +// kill cancels the job's context (terminating the process) and marks it exited if +// it was still running. It reports whether the job was running when called. +func (j *BashJob) kill() bool { + j.mu.Lock() + running := j.status == BashRunning + j.mu.Unlock() + if j.cancel != nil { + j.cancel() + } + return running +} + +// readNew returns the output accumulated since the last read and advances the +// cursor, so successive bash_output calls stream the command's output without +// repeating what was already seen. +func (j *BashJob) readNew() string { + j.mu.Lock() + defer j.mu.Unlock() + all := j.buf.Bytes() + if j.cursor > len(all) { + j.cursor = len(all) + } + out := string(all[j.cursor:]) + j.cursor = len(all) + return out +} + +// snapshot returns the job's current status fields for reporting without exposing +// the mutex-guarded internals. +func (j *BashJob) snapshot() (status BashJobStatus, exitCode int, errMsg string) { + j.mu.Lock() + defer j.mu.Unlock() + return j.status, j.exitCode, j.errMsg +} + +// BashJobStore holds a session's background bash jobs. A single store is shared by +// the bash, bash_output and kill_bash tools so a job launched by one is visible to +// the others. It is safe for concurrent use. +type BashJobStore struct { + mu sync.Mutex + jobs map[string]*BashJob + seq int +} + +// NewBashJobStore returns an empty store. +func NewBashJobStore() *BashJobStore { + return &BashJobStore{jobs: map[string]*BashJob{}} +} + +// create registers a new running job for command with its cancel func, assigning +// a readable sequential id. +func (s *BashJobStore) create(command string, cancel context.CancelFunc) *BashJob { + s.mu.Lock() + defer s.mu.Unlock() + s.seq++ + job := &BashJob{ + ID: fmt.Sprintf("bash_%d", s.seq), + Command: command, + StartedAt: time.Now(), + status: BashRunning, + cancel: cancel, + } + s.jobs[job.ID] = job + return job +} + +// Get returns the job with the given id, or (nil, false). +func (s *BashJobStore) Get(id string) (*BashJob, bool) { + s.mu.Lock() + defer s.mu.Unlock() + j, ok := s.jobs[id] + return j, ok +} + +// List returns the jobs in creation order. +func (s *BashJobStore) List() []*BashJob { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]*BashJob, 0, len(s.jobs)) + for i := 1; i <= s.seq; i++ { + if j, ok := s.jobs[fmt.Sprintf("bash_%d", i)]; ok { + out = append(out, j) + } + } + return out +} + +// KillAll cancels every still-running job. It is intended for session shutdown so +// background processes are not orphaned. +func (s *BashJobStore) KillAll() { + for _, j := range s.List() { + j.kill() + } +} diff --git a/pigo/internal/agenttool/bash_background_test.go b/pigo/internal/agenttool/bash_background_test.go new file mode 100644 index 0000000..c7a352f --- /dev/null +++ b/pigo/internal/agenttool/bash_background_test.go @@ -0,0 +1,134 @@ +// Tests for background bash execution: launching a detached job, draining its +// output incrementally via bash_output, and terminating a long-running job with +// kill_bash. These exercise the shared BashJobStore wiring end to end. +package agenttool + +import ( + "context" + "encoding/json" + "runtime" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func runBGTool(t *testing.T, tool agentcore.AgentTool, args map[string]any) (agentcore.AgentToolResult, error) { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + return tool.Execute(context.Background(), "call-1", raw, nil) +} + +// A background command returns immediately with a bash_id, then bash_output +// drains its output and reports it exited. +func TestBashBackgroundRunAndOutput(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + jobs := NewBashJobStore() + bash := &BashTool{Jobs: jobs} + out := &BashOutputTool{Jobs: jobs} + + res, gerr := runBGTool(t, bash, map[string]any{"command": "echo bg-hello", "run_in_background": true}) + if gerr != nil { + t.Fatalf("unexpected error: %v", gerr) + } + details, ok := res.Details.(map[string]any) + if !ok || details["background"] != true { + t.Fatalf("expected background details, got %+v", res.Details) + } + id, _ := details["bash_id"].(string) + if id == "" { + t.Fatalf("no bash_id returned") + } + + // Poll bash_output until the job exits and produced its line. + var text string + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + r, err := runBGTool(t, out, map[string]any{"bash_id": id}) + if err != nil { + t.Fatalf("bash_output: %v", err) + } + text += resultText(r) + d, _ := r.Details.(map[string]any) + if d["status"] == string(BashExited) { + break + } + time.Sleep(20 * time.Millisecond) + } + if !strings.Contains(text, "bg-hello") { + t.Errorf("output = %q, want to contain bg-hello", text) + } + if !strings.Contains(text, "exited") { + t.Errorf("status never reported exited: %q", text) + } +} + +// kill_bash terminates a long-running background job and it stops running. +func TestBashBackgroundKill(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + jobs := NewBashJobStore() + bash := &BashTool{Jobs: jobs} + kill := &BashKillTool{Jobs: jobs} + out := &BashOutputTool{Jobs: jobs} + + res, gerr := runBGTool(t, bash, map[string]any{"command": "sleep 30", "run_in_background": true}) + if gerr != nil { + t.Fatalf("unexpected error: %v", gerr) + } + id := res.Details.(map[string]any)["bash_id"].(string) + + kr, err := runBGTool(t, kill, map[string]any{"bash_id": id}) + if err != nil { + t.Fatalf("kill_bash: %v", err) + } + if killed, _ := kr.Details.(map[string]any)["killed"].(bool); !killed { + t.Errorf("expected killed=true, got %+v", kr.Details) + } + + // After the kill, the job should report exited within a short window. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + r, _ := runBGTool(t, out, map[string]any{"bash_id": id}) + if r.Details.(map[string]any)["status"] == string(BashExited) { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Errorf("job still running after kill") +} + +// bash_output and kill_bash report a clear error for an unknown id. +func TestBashControlUnknownID(t *testing.T) { + jobs := NewBashJobStore() + out := &BashOutputTool{Jobs: jobs} + kill := &BashKillTool{Jobs: jobs} + + r, _ := runBGTool(t, out, map[string]any{"bash_id": "bash_99"}) + if !strings.Contains(resultText(r), "no background command") { + t.Errorf("bash_output on unknown id should error, got %q", resultText(r)) + } + r, _ = runBGTool(t, kill, map[string]any{"bash_id": "bash_99"}) + if !strings.Contains(resultText(r), "no background command") { + t.Errorf("kill_bash on unknown id should error, got %q", resultText(r)) + } +} + +// run_in_background without a store wired reports it is unavailable. +func TestBashBackgroundNoStore(t *testing.T) { + bash := &BashTool{} + r, err := runBGTool(t, bash, map[string]any{"command": "echo x", "run_in_background": true}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(resultText(r), "not available") { + t.Errorf("expected unavailable message when no store is wired, got %q", resultText(r)) + } +} diff --git a/pigo/internal/agenttool/bash_control.go b/pigo/internal/agenttool/bash_control.go new file mode 100644 index 0000000..360e882 --- /dev/null +++ b/pigo/internal/agenttool/bash_control.go @@ -0,0 +1,158 @@ +// This file implements the two companion tools for background bash jobs: +// bash_output drains a job's new output and reports its status, and kill_bash +// terminates a running job. Both address a job by the bash_id returned from a +// `bash` call with run_in_background=true, sharing the same BashJobStore so a +// job launched by the bash tool is visible here. This mirrors Claude Code's +// BashOutput/KillShell tools. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// BashOutputTool reads the output a background job has produced since the last +// read and reports whether it is still running or has exited. Jobs is the shared +// store the bash tool populates. +type BashOutputTool struct { + Jobs *BashJobStore +} + +// bashOutputArgs is the decoded argument shape for BashOutputTool. +type bashOutputArgs struct { + // BashID is the job handle returned by a background `bash` call. + BashID string `json:"bash_id"` +} + +// Name implements AgentTool. +func (t *BashOutputTool) Name() string { return "bash_output" } + +// Description implements AgentTool. +func (t *BashOutputTool) Description() string { + return "Read new output from a background command started with bash " + + "run_in_background=true, addressed by its bash_id. Returns output " + + "accumulated since the last read plus the command's status (running or " + + "exited, with exit code). Call repeatedly to stream a long job's output." +} + +// Schema implements AgentTool. +func (t *BashOutputTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "bash_id": {"type": "string", "description": "The bash_id returned by a background bash call."} + }, + "required": ["bash_id"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Reading output has no side effects. +func (t *BashOutputTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} + +// Execute implements AgentTool. It returns the job's new output and status. +func (t *BashOutputTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[bashOutputArgs](args, "bash_output") + if bad != nil { + return *bad, nil + } + if t.Jobs == nil { + return errorResult("bash_output: background jobs are not available in this environment"), nil + } + job, ok := t.Jobs.Get(a.BashID) + if !ok { + return errorResult(fmt.Sprintf("bash_output: no background command with id %q", a.BashID)), nil + } + + out := truncateBashOutput(job.readNew()) + status, exitCode, errMsg := job.snapshot() + + var statusLine string + if status == BashRunning { + statusLine = fmt.Sprintf("[%s: running]", a.BashID) + } else if errMsg != "" && exitCode != 0 { + statusLine = fmt.Sprintf("[%s: exited code %d: %s]", a.BashID, exitCode, errMsg) + } else { + statusLine = fmt.Sprintf("[%s: exited code %d]", a.BashID, exitCode) + } + + text := statusLine + if out != "" { + text = out + "\n" + statusLine + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + Details: map[string]any{"bash_id": a.BashID, "status": string(status), "exitCode": exitCode}, + }, nil +} + +// BashKillTool terminates a running background job by canceling its context. +// Jobs is the shared store the bash tool populates. +type BashKillTool struct { + Jobs *BashJobStore +} + +// bashKillArgs is the decoded argument shape for BashKillTool. +type bashKillArgs struct { + // BashID is the job handle returned by a background `bash` call. + BashID string `json:"bash_id"` +} + +// Name implements AgentTool. +func (t *BashKillTool) Name() string { return "kill_bash" } + +// Description implements AgentTool. +func (t *BashKillTool) Description() string { + return "Terminate a background command started with bash " + + "run_in_background=true, addressed by its bash_id. The command's " + + "process is killed; already-exited jobs report that they were not running." +} + +// Schema implements AgentTool. +func (t *BashKillTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "bash_id": {"type": "string", "description": "The bash_id returned by a background bash call."} + }, + "required": ["bash_id"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Killing a process is a side effect. +func (t *BashKillTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// Execute implements AgentTool. It kills the job and reports whether it had been +// running. +func (t *BashKillTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[bashKillArgs](args, "kill_bash") + if bad != nil { + return *bad, nil + } + if t.Jobs == nil { + return errorResult("kill_bash: background jobs are not available in this environment"), nil + } + job, ok := t.Jobs.Get(a.BashID) + if !ok { + return errorResult(fmt.Sprintf("kill_bash: no background command with id %q", a.BashID)), nil + } + + if job.kill() { + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(fmt.Sprintf("killed background command %s", a.BashID))}, + Details: map[string]any{"bash_id": a.BashID, "killed": true}, + }, nil + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(fmt.Sprintf("background command %s was not running", a.BashID))}, + Details: map[string]any{"bash_id": a.BashID, "killed": false}, + }, nil +} diff --git a/pigo/internal/agenttool/bash_tool.go b/pigo/internal/agenttool/bash_tool.go new file mode 100644 index 0000000..3c7abc0 --- /dev/null +++ b/pigo/internal/agenttool/bash_tool.go @@ -0,0 +1,328 @@ +// This file implements the bash tool (US-018): run a shell command, streaming +// stdout/stderr back as tool_execution_update partials, honoring a timeout and +// context cancellation (which kills the child process group). A non-zero exit +// is surfaced as an error (isError) whose message carries the captured output. +package agenttool + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os/exec" + "runtime" + "sync" + "time" + "unicode/utf8" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// bashDefaultTimeout bounds a command that does not specify one. +const bashDefaultTimeout = 2 * time.Minute + +// bashMaxTimeout caps any requested timeout. +const bashMaxTimeout = 10 * time.Minute + +// bashMaxOutputBytes caps how many bytes of combined stdout/stderr the bash tool +// returns to the model. A single command can emit megabytes (build logs, a big +// cat), which — unlike the timeout cap — would otherwise flow into context whole +// and blow the window. Output past this size is truncated to a head + tail +// preview (see truncateBashOutput), mirroring search's searchMaxResults/"[truncated +// …]" convention. This is the tool's own inner cap; a later executor-layer budget +// may impose a stricter outer limit. +const bashMaxOutputBytes = 30_000 + +// truncateBashOutput caps s at bashMaxOutputBytes using the shared +// truncateToBudget idiom (head + "[truncated N bytes]" marker + tail, cut on +// UTF-8 rune boundaries). It is the bash tool's own inner cap; the executor +// layer applies a separate, uniform outer budget afterward. +func truncateBashOutput(s string) string { + return truncateToBudget(s, bashMaxOutputBytes) +} + +// trimUTF8Prefix drops trailing bytes of s that form an incomplete rune, so the +// returned prefix ends on a rune boundary. +func trimUTF8Prefix(s string) string { + for len(s) > 0 { + if r, size := utf8.DecodeLastRuneInString(s); r == utf8.RuneError && size <= 1 { + s = s[:len(s)-1] + continue + } + break + } + return s +} + +// trimUTF8Suffix drops leading bytes of s that form an incomplete rune, so the +// returned suffix starts on a rune boundary. +func trimUTF8Suffix(s string) string { + for len(s) > 0 { + if r, size := utf8.DecodeRuneInString(s); r == utf8.RuneError && size <= 1 { + s = s[1:] + continue + } + break + } + return s +} + +// BashTool runs shell commands. Dir bounds the working directory (empty = the +// process CWD). Shell selects the interpreter (empty = "bash -c"). +type BashTool struct { + // Dir is the working directory for commands. Empty uses the process CWD. + Dir string + // Shell is the interpreter path. Empty defaults to "bash". + Shell string + // Jobs holds background jobs launched with run_in_background. When nil, + // run_in_background is rejected (the front-end did not wire a store). + Jobs *BashJobStore +} + +// bashToolArgs is the decoded argument shape for BashTool. +type bashToolArgs struct { + // Command is the shell command line to run. + Command string `json:"command"` + // TimeoutMs optionally overrides the default timeout (milliseconds). + TimeoutMs int `json:"timeout_ms,omitempty"` + // RunInBackground detaches the command from the turn: it keeps running after + // Execute returns, and its output is drained later via bash_output. A + // background command has no default timeout (so dev servers/watchers run + // indefinitely); timeout_ms still caps it if given. + RunInBackground bool `json:"run_in_background,omitempty"` +} + +// Name implements AgentTool. +func (t *BashTool) Name() string { return "bash" } + +// Description implements AgentTool. +func (t *BashTool) Description() string { + return "Run a shell command, streaming stdout/stderr. Supports a timeout " + + "and cancellation. A non-zero exit code is reported as an error. " + + "Set run_in_background=true for long-running commands (dev servers, " + + "watchers): it returns immediately with a bash_id you drain with " + + "bash_output and stop with kill_bash. " + + "On Windows the command runs under bash if available (Git Bash/WSL), " + + "else PowerShell, else cmd — prefer portable commands." +} + +// Schema implements AgentTool. +func (t *BashTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "command": {"type": "string", "description": "Shell command line to run."}, + "timeout_ms": {"type": "integer", "description": "Timeout in milliseconds (capped at 10 minutes). Ignored in background unless set.", "minimum": 0}, + "run_in_background": {"type": "boolean", "description": "Run detached and return immediately with a bash_id; drain output with bash_output, stop with kill_bash."} + }, + "required": ["command"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Commands can have side effects → sequential. +func (t *BashTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// shellLookPath resolves a program on PATH. It is a package var so tests can +// simulate a Windows box with or without bash installed. +var shellLookPath = exec.LookPath + +// resolveShell picks the interpreter and the flag that makes it read the command +// from the next argument. An explicit shell (BashTool.Shell) is always honored as +// a POSIX-style " -c ". +// +// On Windows with no explicit shell, the naive "bash -c" hardcode fails on stock +// machines that have no bash on PATH — the model then retries bash blindly and +// every call errors (issue #518). So we prefer a real bash when one is present +// (Git Bash / WSL / MSYS), since commands are authored in bash syntax, and fall +// back to PowerShell, then cmd, so a command still runs on a bare Windows box. +func resolveShell(explicit, goos string, lookPath func(string) (string, error)) (shell, flag string) { + if explicit != "" { + return explicit, "-c" + } + if goos == "windows" { + if p, err := lookPath("bash"); err == nil { + return p, "-c" + } + if p, err := lookPath("powershell"); err == nil { + return p, "-Command" + } + return "cmd", "/C" + } + return "bash", "-c" +} + +// streamWriter forwards each written chunk to onUpdate as a growing partial +// result while accumulating the full output. It is safe for concurrent use so +// stdout and stderr can share the same combined buffer. +type streamWriter struct { + mu *sync.Mutex + buf *bytes.Buffer + onUpdate agentcore.ToolUpdateFunc +} + +func (w streamWriter) Write(p []byte) (int, error) { + w.mu.Lock() + w.buf.Write(p) + snapshot := w.buf.String() + w.mu.Unlock() + if w.onUpdate != nil { + w.onUpdate(agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(snapshot)}}) + } + return len(p), nil +} + +// Execute implements AgentTool. It streams combined stdout/stderr via onUpdate, +// enforces a timeout, and kills the process on context cancellation. A non-zero +// exit returns a Go error (→ isError) carrying the exit code and output. +func (t *BashTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[bashToolArgs](args, "bash") + if bad != nil { + return *bad, nil + } + if a.Command == "" { + return errorResult("bash: command is required"), nil + } + + if a.RunInBackground { + return t.startBackground(a) + } + + timeout := bashDefaultTimeout + if a.TimeoutMs > 0 { + timeout = time.Duration(a.TimeoutMs) * time.Millisecond + } + if timeout > bashMaxTimeout { + timeout = bashMaxTimeout + } + + runCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + shell, flag := resolveShell(t.Shell, runtime.GOOS, shellLookPath) + cmd := exec.CommandContext(runCtx, shell, flag, a.Command) + if t.Dir != "" { + cmd.Dir = t.Dir + } + + var mu sync.Mutex + var combined bytes.Buffer + sw := streamWriter{mu: &mu, buf: &combined, onUpdate: onUpdate} + cmd.Stdout = sw + cmd.Stderr = sw + + err := cmd.Run() + + mu.Lock() + output := combined.String() + mu.Unlock() + + // Cap the output before it enters any ToolResult / error message, so a single + // command's huge output cannot blow the model's context. Truncation keeps a + // head + tail preview with a "[truncated N bytes]" marker in the middle. + output = truncateBashOutput(output) + + // Context cancellation / timeout takes precedence in the message. + if runCtx.Err() == context.DeadlineExceeded { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(output)}}, + fmt.Errorf("bash: command timed out after %s\n%s", timeout, output) + } + if ctx.Err() == context.Canceled { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(output)}}, + fmt.Errorf("bash: command canceled\n%s", output) + } + + // A missing interpreter (no bash/powershell/cmd on PATH) surfaces as an + // *exec.Error before the command ever runs. Report it with actionable + // guidance instead of a bare "code -1", so the model stops retrying blindly. + var execErr *exec.Error + if errors.As(err, &execErr) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(output)}}, + fmt.Errorf("bash: could not start shell %q: %v. On Windows install Git Bash or WSL (or configure a shell); commands are bash syntax", shell, execErr.Err) + } + + if err != nil { + exitCode := -1 + var ee *exec.ExitError + if errors.As(err, &ee) { + exitCode = ee.ExitCode() + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(output)}, + Details: map[string]any{"exitCode": exitCode}, + }, + fmt.Errorf("bash: command exited with code %d\n%s", exitCode, output) + } + + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(output)}, + Details: map[string]any{"exitCode": 0}, + }, nil +} + +// startBackground launches the command detached from the turn context and +// returns immediately with a bash_id. The job runs under its own cancelable +// context (rooted at context.Background(), not the turn ctx which is canceled +// when the turn ends), so it survives past Execute. A background command has no +// default timeout — a dev server or watcher is expected to run indefinitely — +// but an explicit timeout_ms still caps it. Its combined output accumulates in +// the job's buffer for bash_output to drain; kill_bash cancels its context. +func (t *BashTool) startBackground(a bashToolArgs) (agentcore.AgentToolResult, error) { + if t.Jobs == nil { + return errorResult("bash: run_in_background is not available in this environment"), nil + } + + var jobCtx context.Context + var cancel context.CancelFunc + if a.TimeoutMs > 0 { + timeout := time.Duration(a.TimeoutMs) * time.Millisecond + if timeout > bashMaxTimeout { + timeout = bashMaxTimeout + } + jobCtx, cancel = context.WithTimeout(context.Background(), timeout) + } else { + jobCtx, cancel = context.WithCancel(context.Background()) + } + + shell, flag := resolveShell(t.Shell, runtime.GOOS, shellLookPath) + cmd := exec.CommandContext(jobCtx, shell, flag, a.Command) + if t.Dir != "" { + cmd.Dir = t.Dir + } + + job := t.Jobs.create(a.Command, cancel) + w := job.writer() + cmd.Stdout = w + cmd.Stderr = w + + if err := cmd.Start(); err != nil { + cancel() + job.finish(-1, err.Error()) + return errorResult(fmt.Sprintf("bash: could not start background command: %v", err)), nil + } + + go func() { + err := cmd.Wait() + cancel() + exitCode := 0 + errMsg := "" + if err != nil { + exitCode = -1 + var ee *exec.ExitError + if errors.As(err, &ee) { + exitCode = ee.ExitCode() + } + errMsg = err.Error() + } + job.finish(exitCode, errMsg) + }() + + msg := fmt.Sprintf("started background command %s: %s\nuse bash_output %q to read its output, kill_bash %q to stop it", job.ID, a.Command, job.ID, job.ID) + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + Details: map[string]any{"bash_id": job.ID, "background": true}, + }, nil +} diff --git a/pigo/internal/agenttool/bash_tool_test.go b/pigo/internal/agenttool/bash_tool_test.go new file mode 100644 index 0000000..4e28b3a --- /dev/null +++ b/pigo/internal/agenttool/bash_tool_test.go @@ -0,0 +1,272 @@ +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "runtime" + "strings" + "sync" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func runBash(t *testing.T, tool *BashTool, args map[string]any, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + return tool.Execute(context.Background(), "call-1", raw, onUpdate) +} + +func TestBashToolSuccess(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + res, gerr := runBash(t, tool, map[string]any{"command": "echo hello"}, nil) + if gerr != nil { + t.Fatalf("unexpected go error: %v", gerr) + } + if !strings.Contains(resultText(res), "hello") { + t.Errorf("output = %q, want to contain hello", resultText(res)) + } + details, ok := res.Details.(map[string]any) + if !ok || details["exitCode"] != 0 { + t.Errorf("expected exitCode 0, details = %+v", res.Details) + } +} + +func TestBashToolNonZeroExitIsError(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + res, gerr := runBash(t, tool, map[string]any{"command": "echo oops >&2; exit 3"}, nil) + // A non-zero exit must surface as a Go error so the executor flags isError. + if gerr == nil { + t.Fatalf("expected go error for non-zero exit, got nil") + } + if !strings.Contains(gerr.Error(), "code 3") { + t.Errorf("error = %q, want to mention code 3", gerr.Error()) + } + // The captured output must ride along. + if !strings.Contains(gerr.Error(), "oops") { + t.Errorf("error = %q, want to carry output", gerr.Error()) + } + details, ok := res.Details.(map[string]any) + if !ok || details["exitCode"] != 3 { + t.Errorf("expected exitCode 3, details = %+v", res.Details) + } +} + +func TestBashToolStreaming(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + var mu sync.Mutex + var updates []string + onUpdate := func(r agentcore.AgentToolResult) { + mu.Lock() + updates = append(updates, resultText(r)) + mu.Unlock() + } + _, gerr := runBash(t, tool, map[string]any{"command": "printf 'a'; printf 'b'"}, onUpdate) + if gerr != nil { + t.Fatalf("unexpected go error: %v", gerr) + } + mu.Lock() + defer mu.Unlock() + if len(updates) == 0 { + t.Fatalf("expected streaming updates, got none") + } + // The final partial should be the full accumulated output. + if last := updates[len(updates)-1]; !strings.Contains(last, "ab") { + t.Errorf("final update = %q, want to contain ab", last) + } +} + +func TestBashToolTimeout(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + start := time.Now() + res, gerr := runBash(t, tool, map[string]any{"command": "sleep 5", "timeout_ms": 100}, nil) + if gerr == nil { + t.Fatalf("expected timeout error, got nil") + } + if !strings.Contains(gerr.Error(), "timed out") { + t.Errorf("error = %q, want to mention timed out", gerr.Error()) + } + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Errorf("timeout took too long: %s (process not killed?)", elapsed) + } + _ = res +} + +func TestBashToolCancel(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + ctx, cancel := context.WithCancel(context.Background()) + raw, _ := json.Marshal(map[string]any{"command": "sleep 5"}) + go func() { + time.Sleep(100 * time.Millisecond) + cancel() + }() + start := time.Now() + _, gerr := tool.Execute(ctx, "call-1", raw, nil) + if gerr == nil { + t.Fatalf("expected cancellation error, got nil") + } + if elapsed := time.Since(start); elapsed > 3*time.Second { + t.Errorf("cancel took too long: %s (process not killed?)", elapsed) + } +} + +func TestBashToolMissingCommand(t *testing.T) { + tool := &BashTool{} + res, gerr := runBash(t, tool, map[string]any{"command": ""}, nil) + if gerr != nil { + t.Fatalf("unexpected go error: %v", gerr) + } + if !strings.Contains(resultText(res), "command is required") { + t.Errorf("expected command-required error, got %q", resultText(res)) + } +} + +func TestBashToolMode(t *testing.T) { + tool := &BashTool{} + if tool.Name() != "bash" { + t.Errorf("name = %q", tool.Name()) + } + if tool.ExecutionMode() != agentcore.ToolExecutionSequential { + t.Error("bash should be sequential") + } + var schema map[string]any + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Errorf("schema not valid JSON: %v", err) + } +} + +func TestBashToolSmallOutputNotTruncated(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + res, gerr := runBash(t, tool, map[string]any{"command": "echo hello world"}, nil) + if gerr != nil { + t.Fatalf("unexpected go error: %v", gerr) + } + out := resultText(res) + if strings.Contains(out, "truncated") { + t.Errorf("small output should not be truncated, got %q", out) + } + if strings.TrimSpace(out) != "hello world" { + t.Errorf("output = %q, want %q", out, "hello world") + } +} + +func TestBashToolLargeOutputTruncatedHeadTail(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on windows") + } + tool := &BashTool{} + // Emit a marker at the very start and very end, with a large filler between, + // so we can prove both the head and the tail survive truncation. + total := bashMaxOutputBytes * 3 + filler := bashMaxOutputBytes // bytes of 'x' between the two markers + cmd := fmt.Sprintf("printf 'HEADMARK'; head -c %d /dev/zero | tr '\\0' 'x'; printf 'TAILMARK'", filler) + _ = total + res, gerr := runBash(t, tool, map[string]any{"command": cmd}, nil) + if gerr != nil { + t.Fatalf("unexpected go error: %v", gerr) + } + out := resultText(res) + if len(out) > bashMaxOutputBytes+128 { + t.Errorf("truncated output too long: %d bytes (cap %d)", len(out), bashMaxOutputBytes) + } + if !strings.HasPrefix(out, "HEADMARK") { + t.Errorf("head not preserved; output starts with %q", out[:min(16, len(out))]) + } + if !strings.HasSuffix(out, "TAILMARK") { + t.Errorf("tail not preserved; output ends with %q", out[max(0, len(out)-16):]) + } + if !strings.Contains(out, "[truncated ") || !strings.Contains(out, " bytes]") { + t.Errorf("missing truncation marker in %q", out) + } +} + +func TestTruncateBashOutputByteCount(t *testing.T) { + // A pure-ASCII input of a known size: the marker's N must equal the exact + // number of middle bytes dropped, i.e. total - head - tail. + total := bashMaxOutputBytes * 2 + in := strings.Repeat("z", total) + out := truncateBashOutput(in) + + half := bashMaxOutputBytes / 2 + // For all-ASCII input no rune-boundary trimming happens, so head/tail are + // each exactly half and N = total - 2*half. + wantRemoved := total - 2*half + wantMarker := fmt.Sprintf("[truncated %d bytes]", wantRemoved) + if !strings.Contains(out, wantMarker) { + t.Errorf("marker = ...%q..., want to contain %q", out, wantMarker) + } + if got := strings.Count(out, "z"); got != 2*half { + t.Errorf("preserved %d content bytes, want %d (head+tail)", got, 2*half) + } + + // Input at or below the cap is returned verbatim. + small := strings.Repeat("b", bashMaxOutputBytes) + if got := truncateBashOutput(small); got != small { + t.Errorf("input at cap should be unchanged") + } +} + +// TestResolveShell covers the platform-aware interpreter selection (issue #518). +// It injects goos + a lookPath stub so every branch runs regardless of the host. +func TestResolveShell(t *testing.T) { + found := func(name string) func(string) (string, error) { + return func(s string) (string, error) { + if s == name { + return `C:\bin\` + s, nil + } + return "", fmt.Errorf("not found") + } + } + none := func(string) (string, error) { return "", fmt.Errorf("not found") } + + tests := []struct { + name string + explicit, goos string + lookPath func(string) (string, error) + wantFlag string + wantShellHas string // substring the resolved shell must contain + }{ + {"explicit honored on windows", "zsh", "windows", none, "-c", "zsh"}, + {"explicit honored on linux", "fish", "linux", none, "-c", "fish"}, + {"non-windows always bash", "", "linux", none, "-c", "bash"}, + {"darwin always bash", "", "darwin", none, "-c", "bash"}, + {"windows with bash", "", "windows", found("bash"), "-c", "bash"}, + {"windows falls back to powershell", "", "windows", found("powershell"), "-Command", "powershell"}, + {"windows falls back to cmd", "", "windows", none, "/C", "cmd"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + shell, flag := resolveShell(tc.explicit, tc.goos, tc.lookPath) + if flag != tc.wantFlag { + t.Errorf("flag = %q, want %q", flag, tc.wantFlag) + } + if !strings.Contains(shell, tc.wantShellHas) { + t.Errorf("shell = %q, want to contain %q", shell, tc.wantShellHas) + } + }) + } +} + diff --git a/pigo/internal/agenttool/batch_executor.go b/pigo/internal/agenttool/batch_executor.go new file mode 100644 index 0000000..cbef614 --- /dev/null +++ b/pigo/internal/agenttool/batch_executor.go @@ -0,0 +1,86 @@ +// This file implements batch tool execution (US-005): a batch of tool calls +// from one assistant message is run either sequentially or in parallel, mirroring +// pi's semantics. +// +// - sequential mode runs each call prepare→execute→finalize in order and stops +// early if the context is aborted. +// - parallel mode preserves ordering by index-backfilling results, running the +// allowed calls in goroutines. (prepare is not separately staged here because +// executeToolCall keeps prepare+execute together per call; ordering is still +// guaranteed by writing each result to its source index.) +// +// The whole batch signals termination only when every finalized result has +// terminate=true, matching pi. +package agenttool + +import ( + "context" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// ForceSequential, when true, makes the whole batch run serially regardless of +// per-tool ExecutionMode. +type BatchConfig struct { + ToolExecutorConfig + ForceSequential bool +} + +// ExecuteToolCalls runs a batch of tool calls belonging to one assistant +// message. It returns the tool-result messages in source order and whether the +// whole batch requests termination (only when every result terminates). +func ExecuteToolCalls(ctx context.Context, cfg BatchConfig, calls []agentcore.AgentToolCall, emit agentcore.EmitFunc) ([]agentcore.ToolResultMessage, bool) { + if len(calls) == 0 { + return nil, false + } + + results := make([]agentcore.ToolResultMessage, len(calls)) + terminates := make([]bool, len(calls)) + + if cfg.ForceSequential || batchRequiresSequential(cfg.Registry, calls) { + for i, call := range calls { + if ctx.Err() != nil { + // Abort: fill the remaining calls with aborted error results so + // every tool call still gets a result message. + for j := i; j < len(calls); j++ { + results[j] = errorToolResult(calls[j], "tool call aborted") + terminates[j] = false + } + break + } + results[i], terminates[i] = executeToolCall(ctx, cfg.ToolExecutorConfig, call, emit) + } + } else { + var wg sync.WaitGroup + for i, call := range calls { + wg.Add(1) + go func(i int, call agentcore.AgentToolCall) { + defer wg.Done() + results[i], terminates[i] = executeToolCall(ctx, cfg.ToolExecutorConfig, call, emit) + }(i, call) + } + wg.Wait() + } + + // Whole batch terminates only when every result terminates (pi semantics). + allTerminate := true + for _, t := range terminates { + if !t { + allTerminate = false + break + } + } + return results, allTerminate +} + +// batchRequiresSequential reports whether any tool in the batch declares +// ExecutionMode sequential, which forces the whole batch to run serially. +func batchRequiresSequential(reg *ToolRegistry, calls []agentcore.AgentToolCall) bool { + for _, call := range calls { + if tool, ok := reg.Get(call.Name); ok && tool.ExecutionMode() == agentcore.ToolExecutionSequential { + return true + } + } + return false +} diff --git a/pigo/internal/agenttool/batch_executor_test.go b/pigo/internal/agenttool/batch_executor_test.go new file mode 100644 index 0000000..1557f5d --- /dev/null +++ b/pigo/internal/agenttool/batch_executor_test.go @@ -0,0 +1,227 @@ +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// registerAll builds a registry containing every tool given. +func registerAll(t *testing.T, tools ...agentcore.AgentTool) *ToolRegistry { + t.Helper() + r := NewToolRegistry() + for _, tool := range tools { + if err := r.Register(tool); err != nil { + t.Fatalf("register %s: %v", tool.Name(), err) + } + } + return r +} + +// echoTool returns its name as text; optionally terminates. +func echoTool(name string, mode agentcore.ToolExecutionMode, terminate bool) execTool { + return execTool{ + name: name, + mode: mode, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + term := terminate + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}, Terminate: &term}, nil + }, + } +} + +func callsFor(names ...string) []agentcore.AgentToolCall { + calls := make([]agentcore.AgentToolCall, len(names)) + for i, n := range names { + calls[i] = agentcore.AgentToolCall{ID: fmt.Sprintf("c%d", i), Name: n, Arguments: json.RawMessage(`{}`)} + } + return calls +} + +// TestBatchParallelPreservesOrder verifies that parallel execution backfills +// results at their source index regardless of completion order. +func TestBatchParallelPreservesOrder(t *testing.T) { + // t0 sleeps longest, t2 shortest — so completion order is reversed, but the + // result slice must still be [t0, t1, t2]. + mk := func(name string, delay time.Duration) execTool { + return execTool{ + name: name, + mode: agentcore.ToolExecutionParallel, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + time.Sleep(delay) + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } + } + reg := registerAll(t, mk("t0", 30*time.Millisecond), mk("t1", 15*time.Millisecond), mk("t2", 1*time.Millisecond)) + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg}} + + results, term := ExecuteToolCalls(context.Background(), cfg, callsFor("t0", "t1", "t2"), nil) + if term { + t.Errorf("no tool terminates; batch must not terminate") + } + want := []string{"t0", "t1", "t2"} + for i, w := range want { + if got := textOf(results[i]); got != w { + t.Errorf("result[%d] = %q, want %q (order not preserved)", i, got, w) + } + } +} + +// TestBatchParallelRunsConcurrently confirms parallel tools overlap in time. +func TestBatchParallelRunsConcurrently(t *testing.T) { + var mu sync.Mutex + running := 0 + maxConcurrent := 0 + block := make(chan struct{}) + var started sync.WaitGroup + started.Add(3) + + mk := func(name string) execTool { + return execTool{ + name: name, + mode: agentcore.ToolExecutionParallel, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + mu.Lock() + running++ + if running > maxConcurrent { + maxConcurrent = running + } + mu.Unlock() + started.Done() + <-block // hold until all have started + mu.Lock() + running-- + mu.Unlock() + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } + } + reg := registerAll(t, mk("a"), mk("b"), mk("c")) + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg}} + + go func() { + started.Wait() + close(block) + }() + ExecuteToolCalls(context.Background(), cfg, callsFor("a", "b", "c"), nil) + + if maxConcurrent < 3 { + t.Errorf("expected 3 concurrent tools, saw max %d", maxConcurrent) + } +} + +// TestBatchSequentialWhenAnyToolSequential forces serial execution and records +// the order tools actually ran in. +func TestBatchSequentialWhenAnyToolSequential(t *testing.T) { + var mu sync.Mutex + var order []string + mk := func(name string, mode agentcore.ToolExecutionMode) execTool { + return execTool{ + name: name, + mode: mode, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + mu.Lock() + order = append(order, name) + mu.Unlock() + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } + } + // "b" is sequential → whole batch runs serially in source order. + reg := registerAll(t, mk("a", agentcore.ToolExecutionParallel), mk("b", agentcore.ToolExecutionSequential), mk("c", agentcore.ToolExecutionParallel)) + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg}} + + ExecuteToolCalls(context.Background(), cfg, callsFor("a", "b", "c"), nil) + want := []string{"a", "b", "c"} + for i, w := range want { + if order[i] != w { + t.Fatalf("sequential order = %v, want %v", order, want) + } + } +} + +// TestBatchForceSequential verifies the global ForceSequential flag serializes +// even all-parallel tools. +func TestBatchForceSequential(t *testing.T) { + var mu sync.Mutex + var order []string + mk := func(name string) execTool { + return execTool{ + name: name, + mode: agentcore.ToolExecutionParallel, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + mu.Lock() + order = append(order, name) + mu.Unlock() + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } + } + reg := registerAll(t, mk("a"), mk("b")) + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg}, ForceSequential: true} + + ExecuteToolCalls(context.Background(), cfg, callsFor("a", "b"), nil) + if len(order) != 2 || order[0] != "a" || order[1] != "b" { + t.Errorf("force-sequential order = %v, want [a b]", order) + } +} + +// TestBatchTerminateOnlyWhenAll checks the whole-batch terminate semantics. +func TestBatchTerminateOnlyWhenAll(t *testing.T) { + // Mixed: one terminates, one does not → batch must NOT terminate. + reg := registerAll(t, echoTool("term", agentcore.ToolExecutionParallel, true), echoTool("noterm", agentcore.ToolExecutionParallel, false)) + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg}} + _, term := ExecuteToolCalls(context.Background(), cfg, callsFor("term", "noterm"), nil) + if term { + t.Errorf("batch with one non-terminating tool must not terminate") + } + + // All terminate → batch terminates. + reg2 := registerAll(t, echoTool("t1", agentcore.ToolExecutionParallel, true), echoTool("t2", agentcore.ToolExecutionParallel, true)) + cfg2 := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg2}} + _, term2 := ExecuteToolCalls(context.Background(), cfg2, callsFor("t1", "t2"), nil) + if !term2 { + t.Errorf("batch with all terminating tools must terminate") + } +} + +// TestBatchSequentialAbort verifies that aborting mid-batch fills the remaining +// calls with aborted error results. +func TestBatchSequentialAbort(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + mk := func(name string) execTool { + return execTool{ + name: name, + mode: agentcore.ToolExecutionSequential, + run: func(c context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + cancel() // abort after the first tool starts + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } + } + reg := registerAll(t, mk("first"), echoTool("second", agentcore.ToolExecutionSequential, false)) + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: reg}} + + results, _ := ExecuteToolCalls(ctx, cfg, callsFor("first", "second"), nil) + if len(results) != 2 { + t.Fatalf("expected 2 results, got %d", len(results)) + } + if !results[1].IsError { + t.Errorf("second (post-abort) result must be an error result") + } +} + +// TestBatchEmpty covers the empty-batch fast path. +func TestBatchEmpty(t *testing.T) { + cfg := BatchConfig{ToolExecutorConfig: ToolExecutorConfig{Registry: NewToolRegistry()}} + results, term := ExecuteToolCalls(context.Background(), cfg, nil, nil) + if results != nil || term { + t.Errorf("empty batch must return (nil, false), got (%v, %v)", results, term) + } +} diff --git a/pigo/internal/agenttool/blackboard_tool.go b/pigo/internal/agenttool/blackboard_tool.go new file mode 100644 index 0000000..e34b615 --- /dev/null +++ b/pigo/internal/agenttool/blackboard_tool.go @@ -0,0 +1,340 @@ +// This file implements the blackboard tool (US-coop): a shared-file-system +// coordination primitive for the pigo coop runner (see coop/). A single pigo +// agent works on the task in $BB: its workspace lives at $BB/workspace, and it +// creates the DONE marker through this tool. The blackboard ($BB) holds +// task.md, the workspace/, and a DONE marker. +// +// Raw file tools cannot safely create the DONE marker: the agent's file tools +// are rooted at its own workspace, and a plain write could race the supervisor. +// So the blackboard is a dedicated tool with three atomic operations: +// +// blackboard action=read [path=...] global snapshot, or one file's contents +// blackboard action=post file=... content=... atomically append a message +// blackboard action=done summary=... atomically create the DONE marker +// +// Append and marker creation are atomic at the OS level (O_APPEND single-write +// for messages, O_CREATE|O_EXCL for DONE), so writes never interleave or +// clobber each other. Every path is validated against the blackboard root to +// forbid traversal. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// maxBlackboardMessageBytes caps a single post. The cap serves two purposes: +// keeps each append a single atomic write, and keeps the tool result from +// ballooning the context. +const maxBlackboardMessageBytes = 32 * 1024 + +// maxBlackboardReadBytes caps how much of one file blackboard read returns. +const maxBlackboardReadBytes = 32 * 1024 + +// BlackboardTool is the task blackboard for the pigo coop runner. It is +// wired into the tool set only when the BB environment variable points at a +// blackboard root (see run.SetupEnv), so ordinary pigo runs never see it. +type BlackboardTool struct { + // Root is the blackboard root directory (the value of $BB). Must be set. + Root string +} + +// Name implements AgentTool. +func (t *BlackboardTool) Name() string { return "blackboard" } + +// Description implements AgentTool. +func (t *BlackboardTool) Description() string { + return "Read and write the task blackboard used by the pigo coop runner (see " + + "coop/). This is the ONLY tool that may touch shared blackboard files " + + "atomically. Actions: read (no path: global snapshot of task.md, workspace " + + "listing, DONE state; with path like \"workspace/exploit.py\": contents of " + + "that one file); post (atomically append a progress note; file must be a " + + "bare .md name under messages/, e.g. \"round--.md\"); done " + + "(atomically create the DONE marker with a final delivery summary, only when " + + "the deliverable is truly complete; fails if DONE already exists). The " + + "blackboard root, current round and your name are available in the " + + "environment as BB, ROUND, NAME." +} + +// Schema implements AgentTool. +func (t *BlackboardTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "action": {"type": "string", "enum": ["read", "post", "done"], "description": "read | post | done"}, + "path": {"type": "string", "description": "For read: a path under the blackboard root to fetch (e.g. messages/round-1-a.md). Omit for the global snapshot."}, + "file": {"type": "string", "description": "For post: bare file name under messages/, must end in .md (e.g. round-1-a.md)."}, + "content": {"type": "string", "description": "For post: the message body (max 32 KiB)."}, + "summary": {"type": "string", "description": "For done: final delivery summary written into DONE."} + }, + "required": ["action"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. It mutates shared files → sequential. +func (t *BlackboardTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +type blackboardArgs struct { + Action string `json:"action"` + Path string `json:"path"` + File string `json:"file"` + Content string `json:"content"` + Summary string `json:"summary"` +} + +// Execute implements AgentTool. +func (t *BlackboardTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[blackboardArgs](args, "blackboard") + if bad != nil { + return *bad, nil + } + if t.Root == "" { + return errorResult("blackboard: no blackboard root configured"), nil + } + switch a.Action { + case "read": + return t.read(a) + case "post": + return t.post(a) + case "done": + return t.done(a) + default: + return errorResult(fmt.Sprintf("blackboard: unknown action %q (want read|post|done)", a.Action)), nil + } +} + +// read returns either one file's contents (when a.Path is set) or a global +// snapshot of the blackboard: task.md, message list, both workspace listings, +// and the DONE state. +func (t *BlackboardTool) read(a blackboardArgs) (agentcore.AgentToolResult, error) { + if strings.TrimSpace(a.Path) != "" { + return t.readFile(a.Path) + } + + var b strings.Builder + if data, err := os.ReadFile(filepath.Join(t.Root, "task.md")); err == nil { + b.WriteString("# task.md\n") + b.WriteString(truncateToBudget(string(data), maxBlackboardReadBytes)) + b.WriteString("\n") + } else { + b.WriteString("# task.md\n\n") + } + + b.WriteString("\n# messages/ (" + strings.Join(readDirNames(t.Root, "messages"), ", ") + ")\n") + for _, name := range readDirListing(t.Root, "messages") { + b.WriteString(" - " + name + "\n") + } + b.WriteString("\n# workspace/ (your workspace)\n") + for _, name := range readDirListing(t.Root, "workspace") { + b.WriteString(" - " + name + "\n") + } + + done := "" + if data, err := os.ReadFile(filepath.Join(t.Root, "DONE")); err == nil { + done = truncateToBudget(string(data), 4096) + } + b.WriteString("\n# DONE\n") + if done == "" { + b.WriteString("\n") + } else { + b.WriteString("EXISTS:\n" + done + "\n") + } + + if env := blackboardEnvSummary(); env != "" { + b.WriteString("\n# environment\n" + env) + } + + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(b.String())}}, nil +} + +// readFile returns the contents of one file under the blackboard root, capped at +// maxBlackboardReadBytes. Paths are validated (no traversal) and must point +// inside the root. +func (t *BlackboardTool) readFile(p string) (agentcore.AgentToolResult, error) { + full, err := t.safePath(p) + if err != nil { + return errorResult("blackboard read: " + err.Error()), nil + } + info, err := os.Stat(full) + if err != nil { + return errorResult(fmt.Sprintf("blackboard read: %s: %v", p, err)), nil + } + if info.IsDir() { + return errorResult(fmt.Sprintf("blackboard read: %s is a directory; only files can be read", p)), nil + } + data, err := os.ReadFile(full) + if err != nil { + return errorResult(fmt.Sprintf("blackboard read: %s: %v", p, err)), nil + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent("# " + p + "\n" + truncateToBudget(string(data), maxBlackboardReadBytes))}, + }, nil +} + +// post atomically appends a message to $BB/messages/. The file name must +// be a bare *.md name (no separators) so a message can never escape the +// messages directory. Appending is a single O_APPEND write → atomic under +// concurrent agents. +func (t *BlackboardTool) post(a blackboardArgs) (agentcore.AgentToolResult, error) { + name := strings.TrimSpace(a.File) + if !validMessageName(name) { + return errorResult("blackboard post: file must be a bare name ending in .md (e.g. \"round-1-a.md\"), no path separators, no \"..\""), nil + } + content := strings.TrimSpace(a.Content) + if content == "" { + return errorResult("blackboard post: content must not be empty"), nil + } + if len(content) > maxBlackboardMessageBytes { + return errorResult(fmt.Sprintf("blackboard post: content too large (%d bytes, max %d)", len(content), maxBlackboardMessageBytes)), nil + } + + dir := filepath.Join(t.Root, "messages") + if err := os.MkdirAll(dir, 0o755); err != nil { + return errorResult("blackboard post: " + err.Error()), nil + } + // O_APPEND + a single Write is atomic on POSIX: concurrent posts never + // interleave bytes. The newline separates this message from the previous one. + f, err := os.OpenFile(filepath.Join(dir, name), os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) + if err != nil { + return errorResult("blackboard post: " + err.Error()), nil + } + _, werr := f.WriteString(content) + cerr := f.Close() + if werr != nil { + return errorResult("blackboard post: write: " + werr.Error()), nil + } + if cerr != nil { + return errorResult("blackboard post: close: " + cerr.Error()), nil + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(fmt.Sprintf("Message appended to messages/%s", name))}, + }, nil +} + +// done atomically creates the DONE marker with a delivery summary. O_CREATE|O_EXCL +// guarantees exactly one agent can create it; a second attempt reports the +// existing marker rather than overwriting it. +func (t *BlackboardTool) done(a blackboardArgs) (agentcore.AgentToolResult, error) { + summary := strings.TrimSpace(a.Summary) + if summary == "" { + return errorResult("blackboard done: summary must not be empty (include the final delivery summary)"), nil + } + header := "Blackboard cooperation DONE\n" + header += "created: " + time.Now().UTC().Format(time.RFC3339) + "\n\n" + content := header + summary + + target := filepath.Join(t.Root, "DONE") + f, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644) + if err != nil { + if os.IsExist(err) { + existing, _ := os.ReadFile(target) + return errorResult("blackboard done: DONE already exists — cooperation already finished:\n" + truncateToBudget(string(existing), 4096)), nil + } + return errorResult("blackboard done: " + err.Error()), nil + } + if _, werr := f.WriteString(content); werr != nil { + f.Close() + return errorResult("blackboard done: write: " + werr.Error()), nil + } + if cerr := f.Close(); cerr != nil { + return errorResult("blackboard done: close: " + cerr.Error()), nil + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent("DONE marker created. Cooperation finished.")}, + Terminate: terminatePtr(), + }, nil +} + +// safePath resolves p against the blackboard root and rejects anything that +// escapes it (.., absolute paths, symlinks are not followed beyond validation of +// the lexical path). +func (t *BlackboardTool) safePath(p string) (string, error) { + if strings.TrimSpace(p) == "" { + return "", fmt.Errorf("empty path") + } + clean := filepath.Clean(p) + if filepath.IsAbs(clean) { + return "", fmt.Errorf("path %q must be relative to the blackboard root", p) + } + rootClean := filepath.Clean(t.Root) + full := filepath.Join(rootClean, clean) + if full != rootClean && !strings.HasPrefix(full, rootClean+string(filepath.Separator)) { + return "", fmt.Errorf("path %q escapes the blackboard root", p) + } + return full, nil +} + +// validMessageName checks a post file name: a bare *.md name with no directory +// components and no "." or ".." tricks. +func validMessageName(name string) bool { + if name == "" || !strings.HasSuffix(name, ".md") { + return false + } + if strings.ContainsAny(name, `/\`) || name == "." || name == ".." { + return false + } + base := strings.TrimSuffix(name, ".md") + if base == "" || strings.HasPrefix(base, ".") || strings.Contains(base, "..") { + return false + } + return true +} + +// readDirNames lists direct child names of a directory under the root (missing +// or unreadable → empty slice). +func readDirNames(root, sub string) []string { + entries, err := os.ReadDir(filepath.Join(root, sub)) + if err != nil { + return nil + } + names := make([]string, 0, len(entries)) + for _, e := range entries { + names = append(names, e.Name()) + } + sort.Strings(names) + return names +} + +// readDirListing lists direct children of a directory under the root with size +// and modification time, sorted by name (missing or unreadable → empty slice). +func readDirListing(root, sub string) []string { + entries, err := os.ReadDir(filepath.Join(root, sub)) + if err != nil { + return nil + } + out := make([]string, 0, len(entries)) + for _, e := range entries { + var size string + if info, err := e.Info(); err == nil && !info.IsDir() { + size = fmt.Sprintf(" (%d bytes, %s)", info.Size(), info.ModTime().UTC().Format("15:04:05")) + } else if err == nil { + size = " (dir)" + } + out = append(out, e.Name()+size) + } + sort.Strings(out) + return out +} + +// blackboardEnvSummary renders the BB/ROUND/NAME environment values so the model +// can address messages and understand the round. Returns "" when none are set. +func blackboardEnvSummary() string { + var b strings.Builder + for _, k := range []string{"BB", "ROUND", "NAME"} { + if v := strings.TrimSpace(os.Getenv(k)); v != "" { + fmt.Fprintf(&b, " %s=%s\n", k, v) + } + } + return strings.TrimRight(b.String(), "\n") +} diff --git a/pigo/internal/agenttool/blackboard_tool_test.go b/pigo/internal/agenttool/blackboard_tool_test.go new file mode 100644 index 0000000..55bc98d --- /dev/null +++ b/pigo/internal/agenttool/blackboard_tool_test.go @@ -0,0 +1,200 @@ +package agenttool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// newTestBlackboard builds a BlackboardTool over a fresh temp dir and a helper +// to run an action and return the text result. +func newTestBlackboard(t *testing.T) (*BlackboardTool, func(args string) string) { + t.Helper() + root := t.TempDir() + tool := &BlackboardTool{Root: root} + run := func(args string) string { + t.Helper() + res, err := tool.Execute(context.Background(), "t1", json.RawMessage(args), nil) + if err != nil { + t.Fatalf("Execute: %v", err) + } + var sb strings.Builder + for _, c := range res.Content { + if txt, ok := c.(agentcore.TextContent); ok { + sb.WriteString(txt.Text) + } + } + return sb.String() + } + return tool, run +} + +func TestBlackboardPostAndRead(t *testing.T) { + tool, run := newTestBlackboard(t) + + got := run(`{"action":"post","file":"round-1-a.md","content":"hello from a"}`) + if !strings.Contains(got, "appended") { + t.Fatalf("post result = %q, want appended confirmation", got) + } + + data, err := os.ReadFile(filepath.Join(tool.Root, "messages", "round-1-a.md")) + if err != nil { + t.Fatalf("message file not written: %v", err) + } + if string(data) != "hello from a" { + t.Fatalf("message content = %q, want %q", data, "hello from a") + } + + // The snapshot lists the message; readFile returns its contents. + snap := run(`{"action":"read"}`) + if !strings.Contains(snap, "round-1-a.md") { + t.Fatalf("snapshot missing the message name:\n%s", snap) + } + if !strings.Contains(snap, "DONE") || !strings.Contains(snap, "not created") { + t.Fatalf("snapshot missing DONE status:\n%s", snap) + } + one := run(`{"action":"read","path":"messages/round-1-a.md"}`) + if !strings.Contains(one, "hello from a") { + t.Fatalf("readFile result missing content:\n%s", one) + } +} + +func TestBlackboardPostRejectsTraversal(t *testing.T) { + tool, run := newTestBlackboard(t) + for _, bad := range []string{ + `{"action":"post","file":"../escape.md","content":"x"}`, + `{"action":"post","file":"a/b.md","content":"x"}`, + `{"action":"post","file":"..","content":"x"}`, + `{"action":"post","file":"notes.txt","content":"x"}`, + `{"action":"post","file":"round.md","content":""}`, + } { + got := run(bad) + if strings.Contains(got, "appended") { + t.Fatalf("post with %s must be rejected, got %q", bad, got) + } + } + if _, err := os.Stat(filepath.Join(tool.Root, "escape.md")); !os.IsNotExist(err) { + t.Fatalf("traversal escaped the root: %v", err) + } +} + +func TestBlackboardReadRejectsTraversal(t *testing.T) { + _, run := newTestBlackboard(t) + for _, bad := range []string{ + `{"action":"read","path":"../outside.md"}`, + `{"action":"read","path":"/etc/passwd"}`, + `{"action":"read","path":"messages"}`, + } { + got := run(bad) + if !strings.Contains(got, "blackboard read:") { + t.Fatalf("read with %s must error, got %q", bad, got) + } + } +} + +func TestBlackboardDoneIsExclusive(t *testing.T) { + tool, run := newTestBlackboard(t) + + got := run(`{"action":"done","summary":"delivered: flag=abc"}`) + if !strings.Contains(got, "DONE marker created") { + t.Fatalf("first done failed: %q", got) + } + data, err := os.ReadFile(filepath.Join(tool.Root, "DONE")) + if err != nil { + t.Fatalf("DONE not written: %v", err) + } + if !strings.Contains(string(data), "flag=abc") { + t.Fatalf("DONE content missing summary: %q", data) + } + + // A second done must report the existing marker, not overwrite it. + got2 := run(`{"action":"done","summary":"another summary"}`) + if !strings.Contains(got2, "already exists") { + t.Fatalf("second done must report existing marker, got %q", got2) + } + data2, _ := os.ReadFile(filepath.Join(tool.Root, "DONE")) + if strings.Contains(string(data2), "another summary") { + t.Fatalf("second done overwrote the marker: %q", data2) + } + + // done with an empty summary is rejected. + if got := run(`{"action":"done","summary":""}`); !strings.Contains(got, "summary must not be empty") { + t.Fatalf("empty-summary done must error, got %q", got) + } +} + +// TestBlackboardPostConcurrentAtomic verifies that parallel posts to the same +// message file never interleave or lose bytes: each message survives whole. +func TestBlackboardPostConcurrentAtomic(t *testing.T) { + tool, _ := newTestBlackboard(t) + const n = 16 + msgs := make([]string, n) + for i := range msgs { + msgs[i] = strings.Repeat("M", 100) + string(rune('A'+i)) + strings.Repeat("N", 100) + } + + var wg sync.WaitGroup + for i := 0; i < n; i++ { + i := i + wg.Add(1) + go func() { + defer wg.Done() + args := json.RawMessage(`{"action":"post","file":"round-1-a.md","content":"` + msgs[i] + `"}`) + if _, err := tool.Execute(context.Background(), "t", args, nil); err != nil { + t.Errorf("concurrent post %d: %v", i, err) + } + }() + } + wg.Wait() + + data, err := os.ReadFile(filepath.Join(tool.Root, "messages", "round-1-a.md")) + if err != nil { + t.Fatalf("read messages: %v", err) + } + got := string(data) + for i, m := range msgs { + if !strings.Contains(got, m) { + t.Fatalf("message %d lost/interleaved in concurrent append:\n%s", i, got) + } + } + // Every message must appear exactly once (no duplication from re-read+write). + for _, m := range msgs { + if strings.Count(got, m) != 1 { + t.Fatalf("message %q appears %d times:\n%s", m, strings.Count(got, m), got) + } + } +} + +func TestBlackboardPostSizeCap(t *testing.T) { + _, run := newTestBlackboard(t) + huge := strings.Repeat("x", maxBlackboardMessageBytes+1) + got := run(`{"action":"post","file":"big.md","content":"` + huge + `"}`) + if !strings.Contains(got, "too large") { + t.Fatalf("oversized post must error, got %q", got) + } +} + +func TestBlackboardUnknownAction(t *testing.T) { + _, run := newTestBlackboard(t) + if got := run(`{"action":"bogus"}`); !strings.Contains(got, "unknown action") { + t.Fatalf("unknown action must error, got %q", got) + } +} + +func TestBlackboardNoRoot(t *testing.T) { + tool := &BlackboardTool{} + res, err := tool.Execute(context.Background(), "t", json.RawMessage(`{"action":"read"}`), nil) + if err != nil { + t.Fatalf("Execute: %v", err) + } + txt := res.Content[0].(agentcore.TextContent).Text + if !strings.Contains(txt, "no blackboard root") { + t.Fatalf("no-root read must error, got %q", txt) + } +} diff --git a/pigo/internal/agenttool/edit_tool.go b/pigo/internal/agenttool/edit_tool.go new file mode 100644 index 0000000..0e08b95 --- /dev/null +++ b/pigo/internal/agenttool/edit_tool.go @@ -0,0 +1,232 @@ +// This file implements the edit tool (US-017): exact string replacement within +// a file. old_string must match exactly; if it is not unique (and replace_all +// is false) the edit is rejected. A unified-style diff of the change is returned +// for the UI to render. Paths resolve against a Root with the same traversal +// guard as the read/write tools. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// EditTool performs exact string replacements in files under Root. +type EditTool struct { + // Root bounds all edits; a path resolving outside Root is rejected. Empty + // Root defaults to the current working directory. + Root string + // ExtraRoots are additional trusted directories an edit may target even though + // they lie outside Root. It exists for the skills directory so the model can + // modify existing skills that live outside the workspace. + ExtraRoots []string + // Snap, when non-nil, records the file's prior content before it is edited so + // the /rewind command can roll the change back. It is shared with the write tool. + Snap *FileSnapshotRecorder +} + +// editToolArgs is the decoded argument shape for EditTool. +type editToolArgs struct { + Path string `json:"path"` + OldString string `json:"old_string"` + NewString string `json:"new_string"` + ReplaceAll bool `json:"replace_all,omitempty"` +} + +// Name implements AgentTool. +func (t *EditTool) Name() string { return "edit" } + +// Description implements AgentTool. +func (t *EditTool) Description() string { + return "Replace an exact string in a file. old_string must be unique unless " + + "replace_all is set. Returns a diff of the change." +} + +// Schema implements AgentTool. +func (t *EditTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "path": {"type": "string", "description": "File path to edit, relative to the workspace root."}, + "old_string": {"type": "string", "description": "Exact text to replace."}, + "new_string": {"type": "string", "description": "Replacement text."}, + "replace_all": {"type": "boolean", "description": "Replace every occurrence instead of requiring a unique match."} + }, + "required": ["path", "old_string", "new_string"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Edits mutate the filesystem → sequential. +func (t *EditTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// resolvePath resolves p against Root (or any ExtraRoots) via the shared +// resolveWithin boundary policy, so every file tool enforces the same +// workspace-escape guard while edits can also reach trusted extra roots. +func (t *EditTool) resolvePath(p string) (string, error) { + if len(t.ExtraRoots) == 0 { + return resolveWithin(t.Root, p) + } + return resolveWithinAny(append([]string{t.Root}, t.ExtraRoots...), p) +} + +// Execute implements AgentTool. Edit failures (no match, non-unique match, +// missing file, out-of-root) are encoded as error results. +func (t *EditTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[editToolArgs](args, "edit") + if bad != nil { + return *bad, nil + } + if a.Path == "" { + return errorResult("edit: path is required"), nil + } + if a.OldString == a.NewString { + return errorResult("edit: old_string and new_string are identical; nothing to change"), nil + } + full, err := t.resolvePath(a.Path) + if err != nil { + return errorResult("edit: " + err.Error()), nil + } + data, err := os.ReadFile(full) + if err != nil { + if os.IsNotExist(err) { + return errorResult(fmt.Sprintf("edit: file %q does not exist", a.Path)), nil + } + return errorResult(fmt.Sprintf("edit: cannot read %q: %v", a.Path, err)), nil + } + original := string(data) + + count := strings.Count(original, a.OldString) + if count == 0 { + return errorResult(fmt.Sprintf("edit: old_string not found in %q", a.Path)), nil + } + if count > 1 && !a.ReplaceAll { + return errorResult(fmt.Sprintf("edit: old_string is not unique in %q (%d matches); provide more context or set replace_all", a.Path, count)), nil + } + + var updated string + if a.ReplaceAll { + updated = strings.ReplaceAll(original, a.OldString, a.NewString) + } else { + updated = strings.Replace(original, a.OldString, a.NewString, 1) + } + + // Snapshot the prior state before mutating so /rewind can restore it. + t.Snap.Record(full) + if err := os.WriteFile(full, []byte(updated), filePerm); err != nil { + return errorResult(fmt.Sprintf("edit: cannot write %q: %v", a.Path, err)), nil + } + + diff := unifiedDiff(a.Path, original, updated) + replaced := 1 + if a.ReplaceAll { + replaced = count + } + msg := fmt.Sprintf("Edited %s (%d replacement(s))\n%s", a.Path, replaced, diff) + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + Details: map[string]any{"path": a.Path, "replacements": replaced, "diff": diff}, + }, nil +} + +// unifiedDiff produces a minimal line-based diff between old and new content. +// It is not a full unified-diff implementation (no hunk coalescing); it emits a +// header plus per-line -/+ markers, which is enough for a UI to render the +// change. Unchanged lines are shown with a leading space for context. +func unifiedDiff(path, oldContent, newContent string) string { + oldLines := splitLinesKeep(oldContent) + newLines := splitLinesKeep(newContent) + + var b strings.Builder + fmt.Fprintf(&b, "--- a/%s\n+++ b/%s\n", path, path) + + // Longest common subsequence over lines drives the -/+ markers. + ops := diffLines(oldLines, newLines) + for _, op := range ops { + switch op.kind { + case diffEqual: + fmt.Fprintf(&b, " %s\n", op.text) + case diffDelete: + fmt.Fprintf(&b, "-%s\n", op.text) + case diffInsert: + fmt.Fprintf(&b, "+%s\n", op.text) + } + } + return b.String() +} + +// splitLinesKeep splits s into lines, dropping a single trailing newline so an +// empty final element is not produced for the common "ends with \n" case. +func splitLinesKeep(s string) []string { + if s == "" { + return nil + } + lines := strings.Split(s, "\n") + if len(lines) > 0 && lines[len(lines)-1] == "" { + lines = lines[:len(lines)-1] + } + return lines +} + +type diffKind int + +const ( + diffEqual diffKind = iota + diffDelete + diffInsert +) + +type diffOp struct { + kind diffKind + text string +} + +// diffLines computes a line diff via a standard LCS dynamic-programming table, +// then backtracks to emit equal/delete/insert ops in order. +func diffLines(a, b []string) []diffOp { + n, m := len(a), len(b) + // lcs[i][j] = length of LCS of a[i:] and b[j:]. + lcs := make([][]int, n+1) + for i := range lcs { + lcs[i] = make([]int, m+1) + } + for i := n - 1; i >= 0; i-- { + for j := m - 1; j >= 0; j-- { + if a[i] == b[j] { + lcs[i][j] = lcs[i+1][j+1] + 1 + } else if lcs[i+1][j] >= lcs[i][j+1] { + lcs[i][j] = lcs[i+1][j] + } else { + lcs[i][j] = lcs[i][j+1] + } + } + } + var ops []diffOp + i, j := 0, 0 + for i < n && j < m { + if a[i] == b[j] { + ops = append(ops, diffOp{diffEqual, a[i]}) + i++ + j++ + } else if lcs[i+1][j] >= lcs[i][j+1] { + ops = append(ops, diffOp{diffDelete, a[i]}) + i++ + } else { + ops = append(ops, diffOp{diffInsert, b[j]}) + j++ + } + } + for ; i < n; i++ { + ops = append(ops, diffOp{diffDelete, a[i]}) + } + for ; j < m; j++ { + ops = append(ops, diffOp{diffInsert, b[j]}) + } + return ops +} diff --git a/pigo/internal/agenttool/edit_tool_test.go b/pigo/internal/agenttool/edit_tool_test.go new file mode 100644 index 0000000..7bfa839 --- /dev/null +++ b/pigo/internal/agenttool/edit_tool_test.go @@ -0,0 +1,177 @@ +package agenttool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func runEdit(t *testing.T, tool *EditTool, args map[string]any) agentcore.AgentToolResult { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + res, gerr := tool.Execute(context.Background(), "call-1", raw, nil) + if gerr != nil { + t.Fatalf("execute returned go error: %v", gerr) + } + return res +} + +func seedFile(t *testing.T, dir, name, content string) string { + t.Helper() + p := filepath.Join(dir, name) + if err := os.WriteFile(p, []byte(content), 0o644); err != nil { + t.Fatalf("seed %s: %v", name, err) + } + return p +} + +func TestEditToolUniqueMatch(t *testing.T) { + dir := t.TempDir() + p := seedFile(t, dir, "f.txt", "alpha\nbeta\ngamma\n") + tool := &EditTool{Root: dir} + res := runEdit(t, tool, map[string]any{"path": "f.txt", "old_string": "beta", "new_string": "BETA"}) + if strings.Contains(resultText(res), "not found") || strings.Contains(resultText(res), "not unique") { + t.Fatalf("unexpected error: %q", resultText(res)) + } + got, _ := os.ReadFile(p) + if string(got) != "alpha\nBETA\ngamma\n" { + t.Errorf("content = %q", got) + } + // Diff present. + if !strings.Contains(resultText(res), "-beta") || !strings.Contains(resultText(res), "+BETA") { + t.Errorf("diff missing markers: %q", resultText(res)) + } +} + +func TestEditToolNonUniqueErrors(t *testing.T) { + dir := t.TempDir() + seedFile(t, dir, "f.txt", "x\nx\nx\n") + tool := &EditTool{Root: dir} + res := runEdit(t, tool, map[string]any{"path": "f.txt", "old_string": "x", "new_string": "y"}) + if !strings.Contains(resultText(res), "not unique") { + t.Errorf("expected non-unique error, got %q", resultText(res)) + } + // File unchanged. + got, _ := os.ReadFile(filepath.Join(dir, "f.txt")) + if string(got) != "x\nx\nx\n" { + t.Errorf("file should be unchanged, got %q", got) + } +} + +func TestEditToolReplaceAll(t *testing.T) { + dir := t.TempDir() + p := seedFile(t, dir, "f.txt", "x\nx\nx\n") + tool := &EditTool{Root: dir} + res := runEdit(t, tool, map[string]any{"path": "f.txt", "old_string": "x", "new_string": "y", "replace_all": true}) + got, _ := os.ReadFile(p) + if string(got) != "y\ny\ny\n" { + t.Errorf("content = %q, want all replaced", got) + } + details, ok := res.Details.(map[string]any) + if !ok || details["replacements"] != 3 { + t.Errorf("expected 3 replacements, details = %+v", res.Details) + } +} + +func TestEditToolNotFound(t *testing.T) { + dir := t.TempDir() + seedFile(t, dir, "f.txt", "hello\n") + tool := &EditTool{Root: dir} + res := runEdit(t, tool, map[string]any{"path": "f.txt", "old_string": "missing", "new_string": "x"}) + if !strings.Contains(resultText(res), "not found") { + t.Errorf("expected not-found error, got %q", resultText(res)) + } +} + +func TestEditToolMissingFile(t *testing.T) { + tool := &EditTool{Root: t.TempDir()} + res := runEdit(t, tool, map[string]any{"path": "nope.txt", "old_string": "a", "new_string": "b"}) + if !strings.Contains(resultText(res), "does not exist") { + t.Errorf("expected does-not-exist, got %q", resultText(res)) + } +} + +func TestEditToolIdenticalStrings(t *testing.T) { + dir := t.TempDir() + seedFile(t, dir, "f.txt", "a\n") + tool := &EditTool{Root: dir} + res := runEdit(t, tool, map[string]any{"path": "f.txt", "old_string": "a", "new_string": "a"}) + if !strings.Contains(resultText(res), "identical") { + t.Errorf("expected identical error, got %q", resultText(res)) + } +} + +func TestEditToolPathTraversal(t *testing.T) { + dir := t.TempDir() + tool := &EditTool{Root: dir} + res := runEdit(t, tool, map[string]any{"path": "../x.txt", "old_string": "a", "new_string": "b"}) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Errorf("expected boundary error, got %q", resultText(res)) + } +} + +func TestEditToolExtraRootsAllowsSkillModification(t *testing.T) { + work := t.TempDir() + skills := t.TempDir() + skillFile := filepath.Join(skills, "weather", "SKILL.md") + if err := os.MkdirAll(filepath.Dir(skillFile), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + if err := os.WriteFile(skillFile, []byte("old body\n"), 0o644); err != nil { + t.Fatalf("seed skill: %v", err) + } + + // Without ExtraRoots the out-of-workspace skill edit is rejected. + bounded := &EditTool{Root: work} + res := runEdit(t, bounded, map[string]any{"path": skillFile, "old_string": "old body", "new_string": "new body"}) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Fatalf("expected boundary rejection without ExtraRoots, got %q", resultText(res)) + } + + // With the skills dir as an extra root the edit applies. + tool := &EditTool{Root: work, ExtraRoots: []string{skills}} + res = runEdit(t, tool, map[string]any{"path": skillFile, "old_string": "old body", "new_string": "new body"}) + if strings.Contains(resultText(res), "outside the workspace root") { + t.Fatalf("edit still blocked with ExtraRoots: %q", resultText(res)) + } + got, _ := os.ReadFile(skillFile) + if !strings.Contains(string(got), "new body") { + t.Fatalf("skill not modified, content = %q", got) + } +} + +func TestEditToolMode(t *testing.T) { + tool := &EditTool{} + if tool.Name() != "edit" { + t.Errorf("name = %q", tool.Name()) + } + if tool.ExecutionMode() != agentcore.ToolExecutionSequential { + t.Error("edit should be sequential") + } + var schema map[string]any + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Errorf("schema not valid JSON: %v", err) + } +} + +func TestUnifiedDiff(t *testing.T) { + diff := unifiedDiff("f.txt", "a\nb\nc\n", "a\nB\nc\n") + if !strings.Contains(diff, "--- a/f.txt") || !strings.Contains(diff, "+++ b/f.txt") { + t.Errorf("missing header: %q", diff) + } + if !strings.Contains(diff, "-b") || !strings.Contains(diff, "+B") { + t.Errorf("missing change lines: %q", diff) + } + // Unchanged context lines carry a leading space. + if !strings.Contains(diff, " a") || !strings.Contains(diff, " c") { + t.Errorf("missing context lines: %q", diff) + } +} diff --git a/pigo/internal/agenttool/file_snapshot.go b/pigo/internal/agenttool/file_snapshot.go new file mode 100644 index 0000000..25b7611 --- /dev/null +++ b/pigo/internal/agenttool/file_snapshot.go @@ -0,0 +1,213 @@ +// This file implements FileSnapshotRecorder, the edit-rewind journal backing the +// /rewind command. Before the write and edit tools mutate a file they call +// Record(absPath), which captures the file's prior content (or notes that it did +// not exist). Snapshots accumulate per turn; Commit groups the turn's snapshots +// into a RestorePoint tagged with the conversation leaf that preceded the turn. +// Restore replays a suffix of the restore points in reverse to roll the working +// tree back to an earlier state, mirroring Claude Code's Esc-Esc rewind. The +// journal is in-memory and scoped to the running session; only pigo's own +// write/edit tools are captured (arbitrary bash edits are not). +package agenttool + +import ( + "fmt" + "os" + "sync" + "time" +) + +// snapshotMaxBytes caps how large a file may be for its prior content to be held +// in the rewind journal. A file above this is still recorded (so rewind knows it +// changed) but its content is not retained, and rewind reports it as skipped +// rather than clobbering it with stale bytes. +const snapshotMaxBytes = 16 * 1024 * 1024 + +// fileSnapshot is the pre-mutation state of a single file: its content before the +// first write/edit of a turn, or a marker that it did not yet exist (so rewind +// deletes it). TooLarge marks a file that exceeded snapshotMaxBytes, whose +// content was not retained. +type fileSnapshot struct { + Path string + Existed bool + TooLarge bool + Content []byte +} + +// RestorePoint is one turn's worth of file snapshots plus the conversation leaf +// that preceded the turn. Rewinding to it restores every file to its Snapshots +// state and moves the active conversation leaf back to LeafID. +type RestorePoint struct { + Seq int + Time time.Time + LeafID string + Label string + Snapshots []fileSnapshot +} + +// FileSnapshotRecorder captures prior file content before write/edit mutations +// and groups it into per-turn RestorePoints. Its methods are safe for concurrent +// use so parallel tool calls within a turn can record without racing. +type FileSnapshotRecorder struct { + mu sync.Mutex + pending map[string]fileSnapshot // absolute path -> first snapshot this turn + order []string // first-touch order within the turn + points []RestorePoint + nextSeq int +} + +// NewFileSnapshotRecorder returns an empty recorder ready to record the first +// turn's mutations. +func NewFileSnapshotRecorder() *FileSnapshotRecorder { + return &FileSnapshotRecorder{pending: map[string]fileSnapshot{}, nextSeq: 1} +} + +// Record captures the current on-disk state of absPath before it is mutated. Only +// the first call for a given path within a turn is retained, so the snapshot +// reflects the state at the turn's start (later mutations in the same turn are +// rolled back to that same baseline). A nil recorder is a no-op, so tools can +// hold an always-safe optional handle. +func (r *FileSnapshotRecorder) Record(absPath string) { + if r == nil { + return + } + r.mu.Lock() + defer r.mu.Unlock() + if _, seen := r.pending[absPath]; seen { + return + } + snap := fileSnapshot{Path: absPath} + info, err := os.Stat(absPath) + switch { + case err != nil: + // Treat any stat error (including not-exist) as "did not exist": rewind will + // delete the file created this turn. + snap.Existed = false + case info.IsDir(): + // A directory is never written by the file tools; skip content capture. + snap.Existed = true + snap.TooLarge = true + case info.Size() > snapshotMaxBytes: + snap.Existed = true + snap.TooLarge = true + default: + data, readErr := os.ReadFile(absPath) + if readErr != nil { + snap.Existed = true + snap.TooLarge = true + } else { + snap.Existed = true + snap.Content = data + } + } + r.pending[absPath] = snap + r.order = append(r.order, absPath) +} + +// Commit closes the current turn: if any files were recorded it appends a +// RestorePoint tagged with leafID (the conversation leaf before the turn) and +// label (a short description, e.g. the user prompt), then clears the pending +// buffer. A turn that mutated no files creates no restore point. It reports +// whether a restore point was created. +func (r *FileSnapshotRecorder) Commit(leafID, label string) bool { + if r == nil { + return false + } + r.mu.Lock() + defer r.mu.Unlock() + if len(r.order) == 0 { + return false + } + snaps := make([]fileSnapshot, 0, len(r.order)) + for _, p := range r.order { + snaps = append(snaps, r.pending[p]) + } + r.points = append(r.points, RestorePoint{ + Seq: r.nextSeq, + Time: time.Now().UTC(), + LeafID: leafID, + Label: label, + Snapshots: snaps, + }) + r.nextSeq++ + r.pending = map[string]fileSnapshot{} + r.order = nil + return true +} + +// Points returns a copy of the committed restore points, oldest first. +func (r *FileSnapshotRecorder) Points() []RestorePoint { + if r == nil { + return nil + } + r.mu.Lock() + defer r.mu.Unlock() + out := make([]RestorePoint, len(r.points)) + copy(out, r.points) + return out +} + +// Restore rolls the working tree back to the state before the restore point at +// index idx (0-based into the Points slice). It replays that point and every +// later point in reverse, restoring each file's prior content (or deleting files +// that did not exist), then drops those points from the journal so the next +// rewind starts from the new tip. It returns the conversation leaf to switch to +// (the target point's LeafID), the list of restored file paths, and any +// non-fatal warnings (e.g. files skipped because they were too large or a +// restore write failed). +func (r *FileSnapshotRecorder) Restore(idx int) (leafID string, restored []string, warnings []string, err error) { + if r == nil { + return "", nil, nil, fmt.Errorf("no restore points") + } + r.mu.Lock() + defer r.mu.Unlock() + if idx < 0 || idx >= len(r.points) { + return "", nil, nil, fmt.Errorf("restore point %d out of range (have %d)", idx+1, len(r.points)) + } + leafID = r.points[idx].LeafID + + // A file touched across several turns must end at its OLDEST (pre-target) + // baseline. Iterate points oldest→newest and keep only the first snapshot seen + // for each path, so the earliest baseline is the one applied. + applied := map[string]bool{} + for i := idx; i < len(r.points); i++ { + for _, s := range r.points[i].Snapshots { + if applied[s.Path] { + continue + } + applied[s.Path] = true + if w := applySnapshot(s); w != "" { + warnings = append(warnings, w) + continue + } + restored = append(restored, s.Path) + } + } + + r.points = r.points[:idx] + if len(r.points) > 0 { + r.nextSeq = r.points[len(r.points)-1].Seq + 1 + } else { + r.nextSeq = 1 + } + return leafID, restored, warnings, nil +} + +// applySnapshot restores one file to its recorded prior state: rewrite the prior +// content, or delete the file if it did not exist before. It returns a warning +// string when the file cannot be safely restored (too large to have retained +// content, or a filesystem error), or "" on success. +func applySnapshot(s fileSnapshot) string { + if s.TooLarge { + return fmt.Sprintf("%s: skipped (too large to snapshot; left unchanged)", s.Path) + } + if !s.Existed { + if err := os.Remove(s.Path); err != nil && !os.IsNotExist(err) { + return fmt.Sprintf("%s: could not delete: %v", s.Path, err) + } + return "" + } + if err := os.WriteFile(s.Path, s.Content, filePerm); err != nil { + return fmt.Sprintf("%s: could not restore: %v", s.Path, err) + } + return "" +} diff --git a/pigo/internal/agenttool/file_snapshot_test.go b/pigo/internal/agenttool/file_snapshot_test.go new file mode 100644 index 0000000..4eab038 --- /dev/null +++ b/pigo/internal/agenttool/file_snapshot_test.go @@ -0,0 +1,152 @@ +// Tests for the file-snapshot rewind journal: per-turn dedup of recorded paths, +// commit grouping (an untouched turn produces no restore point), and restore +// semantics — reverse replay across turns rolls a file to its oldest baseline, +// files that did not exist are deleted, and the journal is truncated to the tip. +package agenttool + +import ( + "os" + "path/filepath" + "testing" +) + +func writeFileT(t *testing.T, path, content string) { + t.Helper() + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatalf("write %s: %v", path, err) + } +} + +func readFileT(t *testing.T, path string) string { + t.Helper() + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read %s: %v", path, err) + } + return string(data) +} + +// A turn's first Record for a path is the baseline; later Records that turn are +// ignored, and Commit groups the turn's files into one point. +func TestRecorderRecordDedupAndCommit(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "a.txt") + writeFileT(t, f, "v0") + + r := NewFileSnapshotRecorder() + r.Record(f) // baseline "v0" + writeFileT(t, f, "v1") // model's first edit + r.Record(f) // second edit same turn — must be ignored + writeFileT(t, f, "v2") + + if !r.Commit("leaf0", "edit a") { + t.Fatal("Commit reported no restore point despite a recorded file") + } + // An untouched turn creates nothing. + if r.Commit("leaf1", "no edits") { + t.Fatal("Commit created a restore point for a turn with no records") + } + points := r.Points() + if len(points) != 1 || len(points[0].Snapshots) != 1 { + t.Fatalf("want 1 point with 1 snapshot, got %+v", points) + } + if got := string(points[0].Snapshots[0].Content); got != "v0" { + t.Errorf("baseline content = %q, want v0", got) + } +} + +// Restoring rolls files back and deletes ones that did not exist before, and +// returns the pre-turn leaf id. +func TestRecorderRestore(t *testing.T) { + dir := t.TempDir() + existing := filepath.Join(dir, "keep.txt") + created := filepath.Join(dir, "new.txt") + writeFileT(t, existing, "orig") + + r := NewFileSnapshotRecorder() + + // Turn 1: edit an existing file. + r.Record(existing) + writeFileT(t, existing, "edited") + r.Commit("leafA", "turn1") + + // Turn 2: create a brand-new file. + r.Record(created) + writeFileT(t, created, "brand new") + r.Commit("leafB", "turn2") + + // Rewind to before turn 1 (index 0): both turns roll back. + leaf, restored, warnings, err := r.Restore(0) + if err != nil { + t.Fatalf("Restore: %v", err) + } + if leaf != "leafA" { + t.Errorf("target leaf = %q, want leafA", leaf) + } + if len(warnings) != 0 { + t.Errorf("unexpected warnings: %v", warnings) + } + if len(restored) != 2 { + t.Errorf("restored %d files, want 2", len(restored)) + } + if got := readFileT(t, existing); got != "orig" { + t.Errorf("existing file = %q, want orig", got) + } + if _, err := os.Stat(created); !os.IsNotExist(err) { + t.Errorf("created file should have been deleted, stat err = %v", err) + } + if len(r.Points()) != 0 { + t.Errorf("journal should be empty after restoring from index 0") + } +} + +// When a file is edited across several turns, restoring to before the earliest +// of them lands it at its oldest baseline (not an intermediate version). +func TestRecorderRestoreOldestBaselineWins(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "a.txt") + writeFileT(t, f, "v0") + + r := NewFileSnapshotRecorder() + r.Record(f) // baseline v0 + writeFileT(t, f, "v1") + r.Commit("leaf0", "t1") + + r.Record(f) // baseline v1 + writeFileT(t, f, "v2") + r.Commit("leaf1", "t2") + + // Rewind to before t2 only (index 1): file returns to v1. + if _, _, _, err := r.Restore(1); err != nil { + t.Fatalf("Restore(1): %v", err) + } + if got := readFileT(t, f); got != "v1" { + t.Errorf("after rewind to before t2, file = %q, want v1", got) + } + // One point remains (t1); rewind it too → v0. + if _, _, _, err := r.Restore(0); err != nil { + t.Fatalf("Restore(0): %v", err) + } + if got := readFileT(t, f); got != "v0" { + t.Errorf("after full rewind, file = %q, want v0", got) + } +} + +func TestRecorderRestoreOutOfRange(t *testing.T) { + r := NewFileSnapshotRecorder() + if _, _, _, err := r.Restore(0); err == nil { + t.Error("Restore on empty journal should error") + } +} + +// A nil recorder is a safe no-op so tools can hold an optional handle. +func TestRecorderNilSafe(t *testing.T) { + var r *FileSnapshotRecorder + r.Record("/nonexistent") + if r.Commit("x", "y") { + t.Error("nil Commit should report no point") + } + if r.Points() != nil { + t.Error("nil Points should be nil") + } +} diff --git a/pigo/internal/agenttool/goal_tool.go b/pigo/internal/agenttool/goal_tool.go new file mode 100644 index 0000000..10a8527 --- /dev/null +++ b/pigo/internal/agenttool/goal_tool.go @@ -0,0 +1,374 @@ +// This file implements the goal state and the two goal-control tools that power +// the /goal command (mirrors pi-goal / Claude Code's goal mode): given a high-level +// objective, the agent runs autonomously — re-prompted turn after turn — until +// it either declares the goal done (goal_complete), hits a true impasse +// (goal_blocked), or a safety guard / token budget stops it. +// +// The tools live here (rather than in the REPL) because they are ordinary +// AgentTools the model invokes, and because the runtime's GoalReminderProvider +// needs to read the same state — mirroring how TodoTool/TodoStore pairs with +// TodoReminderProvider. GoalState is the shared, concurrency-safe handle both +// the tools (which may run in a batch) and the REPL/reminder (which reads it +// each turn) touch. +package agenttool + +import ( + "context" + "encoding/json" + "strings" + "sync" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// GoalStatus is the lifecycle state of the active goal. +type GoalStatus string + +const ( + // GoalIdle means no goal is set (the zero value). + GoalIdle GoalStatus = "" + // GoalActive means the agent is autonomously working toward the goal. + GoalActive GoalStatus = "active" + // GoalPaused means autonomous continuation stopped (a safety guard fired, or + // the user paused it); it can be resumed. + GoalPaused GoalStatus = "paused" + // GoalBlocked means the agent hit a true impasse (goal_blocked was called). + GoalBlocked GoalStatus = "blocked" + // GoalComplete means the agent declared the goal done (goal_complete). + GoalComplete GoalStatus = "complete" + // GoalBudgetLimited means the token budget was exhausted before completion. + GoalBudgetLimited GoalStatus = "budget_limited" +) + +// GoalState holds the current goal for a REPL session. It is safe for concurrent +// use so the goal tools (which may run in a batch) and the REPL/reminder reader +// can touch it without racing. A single state is shared for a session's +// lifetime; /goal clear resets it to the idle zero value. +type GoalState struct { + mu sync.RWMutex + + id string + objective string + summary string // set by goal_complete + blockReason string // set by goal_blocked + status GoalStatus + + iterations int // autonomous continuations issued so far + noProgress int // consecutive settles with no tool activity + + tokenBudget int // 0 = unlimited + tokensUsed int + + startedAt time.Time +} + +// NewGoalState returns an empty (idle) goal state. +func NewGoalState() *GoalState { return &GoalState{} } + +// GoalSnapshot is an immutable copy of the goal state for display/decisions. +type GoalSnapshot struct { + ID string + Objective string + Summary string + BlockReason string + Status GoalStatus + Iterations int + NoProgress int + TokenBudget int + TokensUsed int + StartedAt time.Time +} + +// Start (re)initializes the state for a new objective, moving it to active. It +// resets all counters so a fresh goal never inherits a prior goal's tallies. +func (s *GoalState) Start(id, objective string, tokenBudget int) { + s.mu.Lock() + defer s.mu.Unlock() + s.id = id + s.objective = objective + s.summary = "" + s.blockReason = "" + s.status = GoalActive + s.iterations = 0 + s.noProgress = 0 + s.tokenBudget = tokenBudget + s.tokensUsed = 0 + s.startedAt = time.Now() +} + +// Snapshot returns a copy of the current state, safe to read without a lock. +func (s *GoalState) Snapshot() GoalSnapshot { + s.mu.RLock() + defer s.mu.RUnlock() + return GoalSnapshot{ + ID: s.id, + Objective: s.objective, + Summary: s.summary, + BlockReason: s.blockReason, + Status: s.status, + Iterations: s.iterations, + NoProgress: s.noProgress, + TokenBudget: s.tokenBudget, + TokensUsed: s.tokensUsed, + StartedAt: s.startedAt, + } +} + +// Clear resets the state to idle (no goal). It zeroes the fields individually +// rather than replacing the whole struct so the embedded mutex (currently held) +// is preserved — overwriting it while locked would corrupt the lock. +func (s *GoalState) Clear() { + s.mu.Lock() + defer s.mu.Unlock() + s.id = "" + s.objective = "" + s.summary = "" + s.blockReason = "" + s.status = GoalIdle + s.iterations = 0 + s.noProgress = 0 + s.tokenBudget = 0 + s.tokensUsed = 0 + s.startedAt = time.Time{} +} + +// SetStatus transitions the goal to a new status (used by the REPL when a safety +// guard fires or the user pauses/resumes). +func (s *GoalState) SetStatus(status GoalStatus) { + s.mu.Lock() + defer s.mu.Unlock() + s.status = status +} + +// Resume reactivates a paused or budget-limited goal and clears the transient +// safety-guard counters (iterations, no-progress) that stopped it, so the run +// gets a fresh allowance rather than immediately re-tripping the same guard. The +// token budget is intentionally reset too (tokensUsed → 0): resuming past an +// exhausted budget is an explicit user decision to grant another window. The +// objective and id are preserved. It is a no-op-safe wrapper — callers gate on +// the current status before invoking it. +func (s *GoalState) Resume() { + s.mu.Lock() + defer s.mu.Unlock() + s.status = GoalActive + s.iterations = 0 + s.noProgress = 0 + s.tokensUsed = 0 +} + +// ID returns the current goal id (empty when idle). +func (s *GoalState) ID() string { + s.mu.RLock() + defer s.mu.RUnlock() + return s.id +} + +// RecordIteration increments the autonomous-continuation counter and folds the +// given output-token delta into the running total. hadToolActivity resets the +// no-progress counter when true, else increments it — so a run of tool-free +// turns can trip the no-progress guard. +func (s *GoalState) RecordIteration(outputTokens int, hadToolActivity bool) { + s.mu.Lock() + defer s.mu.Unlock() + s.iterations++ + s.tokensUsed += outputTokens + if hadToolActivity { + s.noProgress = 0 + } else { + s.noProgress++ + } +} + +// MarkComplete records the completion summary and moves the goal to complete. +func (s *GoalState) MarkComplete(summary string) { + s.mu.Lock() + defer s.mu.Unlock() + s.summary = summary + s.status = GoalComplete +} + +// MarkBlocked records the block reason and moves the goal to blocked. +func (s *GoalState) MarkBlocked(reason string) { + s.mu.Lock() + defer s.mu.Unlock() + s.blockReason = reason + s.status = GoalBlocked +} + +// terminate is the shared *bool=true value returned by both goal tools to end +// the run immediately (the loop terminates when every result in a batch has +// Terminate=true; a goal tool is expected to be the sole call in its turn). +func terminatePtr() *bool { b := true; return &b } + +// contradictorySummary reports whether a goal_complete summary plainly claims the +// goal is NOT done — a guard against the model closing a goal it just admitted +// is unfinished. The check is a conservative substring match on well-known +// negative phrasings (English), matching pi-goal's "plainly contradictory +// summary" rejection. +func contradictorySummary(summary string) bool { + lower := strings.ToLower(summary) + for _, bad := range []string{ + "not complete", + "not done", + "incomplete", + "tests still fail", + "tests fail", + "still failing", + "could not", + "couldn't", + "unable to", + "unfinished", + "did not finish", + "failed to", + "cannot complete", + "still fails", + "test failure", + } { + if strings.Contains(lower, bad) || strings.Contains(summary, bad) { + return true + } + } + return false +} + +// GoalCompleteTool lets the model declare the active goal finished. It records +// the summary, moves the state to complete, and terminates the run. +type GoalCompleteTool struct { + // State is the session goal state. Must be non-nil. + State *GoalState +} + +type goalCompleteArgs struct { + Summary string `json:"summary"` +} + +// Name implements AgentTool. +func (t *GoalCompleteTool) Name() string { return "goal_complete" } + +// Description implements AgentTool. +func (t *GoalCompleteTool) Description() string { + return "Declare the current goal COMPLETE. Call this ONLY after verifying, " + + "requirement by requirement, that the objective is fully met — treat the " + + "working tree, tests, and actual runtime behavior as authoritative, not " + + "the prior conversation. Provide a concise summary of what was accomplished. " + + "Do NOT call this if any requirement is unmet, tests fail, or work remains; " + + "use goal_blocked for a true impasse instead. Calling this ends the run." +} + +// Schema implements AgentTool. +func (t *GoalCompleteTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "summary": {"type": "string", "description": "Concise summary of what was accomplished to satisfy the goal."} + }, + "required": ["summary"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. It mutates shared goal state → sequential. +func (t *GoalCompleteTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// Execute implements AgentTool. It validates the summary (non-empty and not +// plainly contradictory), records completion, and terminates the run. Invalid +// input degrades to an error result (not a Go error) so the model can retry. +func (t *GoalCompleteTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[goalCompleteArgs](args, "goal_complete") + if bad != nil { + return *bad, nil + } + if t.State == nil { + return errorResult("goal_complete: no active goal"), nil + } + summary := strings.TrimSpace(a.Summary) + if summary == "" { + return errorResult("goal_complete: summary must not be empty"), nil + } + if contradictorySummary(summary) { + return errorResult("goal_complete: summary indicates the goal is NOT complete; " + + "keep working, or call goal_blocked with evidence if truly stuck"), nil + } + snap := t.State.Snapshot() + if snap.Status == GoalIdle { + return errorResult("goal_complete: no active goal"), nil + } + t.State.MarkComplete(summary) + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent("Goal marked complete: " + summary)}, + Terminate: terminatePtr(), + }, nil +} + +// GoalBlockedTool lets the model report a true impasse it cannot work around. It +// records the reason, moves the state to blocked, and terminates the run. +type GoalBlockedTool struct { + // State is the session goal state. Must be non-nil. + State *GoalState +} + +type goalBlockedArgs struct { + Reason string `json:"reason"` + Evidence string `json:"evidence"` +} + +// Name implements AgentTool. +func (t *GoalBlockedTool) Name() string { return "goal_blocked" } + +// Description implements AgentTool. +func (t *GoalBlockedTool) Description() string { + return "Report that the current goal is BLOCKED by a true impasse you cannot " + + "resolve (e.g. missing credentials, an external dependency you cannot " + + "install, contradictory requirements). Provide a concrete reason and the " + + "evidence that establishes the blocker. Use this only as a last resort — " + + "prefer trying a different approach first. Calling this ends the run." +} + +// Schema implements AgentTool. +func (t *GoalBlockedTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "reason": {"type": "string", "description": "Concise statement of what blocks the goal."}, + "evidence": {"type": "string", "description": "Concrete evidence establishing the blocker (error output, missing file, etc.)."} + }, + "required": ["reason", "evidence"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. It mutates shared goal state → sequential. +func (t *GoalBlockedTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// Execute implements AgentTool. It validates the reason/evidence, records the +// block, and terminates the run. +func (t *GoalBlockedTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[goalBlockedArgs](args, "goal_blocked") + if bad != nil { + return *bad, nil + } + if t.State == nil { + return errorResult("goal_blocked: no active goal"), nil + } + reason := strings.TrimSpace(a.Reason) + if reason == "" { + return errorResult("goal_blocked: reason must not be empty"), nil + } + if strings.TrimSpace(a.Evidence) == "" { + return errorResult("goal_blocked: evidence must not be empty"), nil + } + snap := t.State.Snapshot() + if snap.Status == GoalIdle { + return errorResult("goal_blocked: no active goal"), nil + } + t.State.MarkBlocked(reason) + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent("Goal marked blocked: " + reason)}, + Terminate: terminatePtr(), + }, nil +} diff --git a/pigo/internal/agenttool/goal_tool_test.go b/pigo/internal/agenttool/goal_tool_test.go new file mode 100644 index 0000000..c8d392c --- /dev/null +++ b/pigo/internal/agenttool/goal_tool_test.go @@ -0,0 +1,188 @@ +// Tests for the goal tools (mirrors pi-goal): goal_complete validation (empty and +// contradictory summaries are rejected, a valid summary marks complete and +// terminates the run), goal_blocked validation, and GoalState counter/lifecycle +// behavior. Mirrors todo_tool_test.go's structure. +package agenttool + +import ( + "context" + "encoding/json" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// execGoalComplete runs the goal_complete tool with the given JSON args. +func execGoalComplete(t *testing.T, tool *GoalCompleteTool, args string) agentcore.AgentToolResult { + t.Helper() + res, err := tool.Execute(context.Background(), "call-1", json.RawMessage(args), nil) + if err != nil { + t.Fatalf("Execute returned Go error: %v", err) + } + return res +} + +func isErrorResult(res agentcore.AgentToolResult) bool { + // An error result carries no Terminate and its text is the error message; + // the tools return errorResult(...) which has Terminate=nil. + return res.Terminate == nil +} + +func TestGoalToolsRegister(t *testing.T) { + reg := NewToolRegistry() + st := NewGoalState() + if err := reg.Register(&GoalCompleteTool{State: st}); err != nil { + t.Fatalf("Register goal_complete: %v", err) + } + if err := reg.Register(&GoalBlockedTool{State: st}); err != nil { + t.Fatalf("Register goal_blocked: %v", err) + } + if _, ok := reg.Get("goal_complete"); !ok { + t.Fatal("goal_complete not found after Register") + } + if _, ok := reg.Get("goal_blocked"); !ok { + t.Fatal("goal_blocked not found after Register") + } +} + +func TestGoalCompleteValidSummary(t *testing.T) { + st := NewGoalState() + st.Start("g1", "do the thing", 0) + tool := &GoalCompleteTool{State: st} + + res := execGoalComplete(t, tool, `{"summary":"created hello.txt with the requested contents"}`) + if res.Terminate == nil || !*res.Terminate { + t.Fatalf("expected Terminate=true, got %v", res.Terminate) + } + if snap := st.Snapshot(); snap.Status != GoalComplete { + t.Errorf("status = %q, want complete", snap.Status) + } +} + +func TestGoalCompleteRejectsEmpty(t *testing.T) { + st := NewGoalState() + st.Start("g1", "do the thing", 0) + tool := &GoalCompleteTool{State: st} + + res := execGoalComplete(t, tool, `{"summary":" "}`) + if !isErrorResult(res) { + t.Fatal("expected error result for empty summary") + } + if snap := st.Snapshot(); snap.Status != GoalActive { + t.Errorf("status = %q, want still active after rejected summary", snap.Status) + } +} + +func TestGoalCompleteRejectsContradictory(t *testing.T) { + st := NewGoalState() + st.Start("g1", "do the thing", 0) + tool := &GoalCompleteTool{State: st} + + for _, summary := range []string{ + `{"summary":"the goal is not complete yet"}`, + `{"summary":"tests still fail but I stopped"}`, + `{"summary":"the task is unfinished"}`, + } { + res := execGoalComplete(t, tool, summary) + if !isErrorResult(res) { + t.Fatalf("expected error result for contradictory summary %s", summary) + } + } + if snap := st.Snapshot(); snap.Status != GoalActive { + t.Errorf("status = %q, want still active", snap.Status) + } +} + +func TestGoalBlockedRecordsReason(t *testing.T) { + st := NewGoalState() + st.Start("g1", "do the thing", 0) + tool := &GoalBlockedTool{State: st} + + res, err := tool.Execute(context.Background(), "c1", + json.RawMessage(`{"reason":"missing API key","evidence":"env AUTH_TOKEN is empty"}`), nil) + if err != nil { + t.Fatalf("Execute: %v", err) + } + if res.Terminate == nil || !*res.Terminate { + t.Fatalf("expected Terminate=true, got %v", res.Terminate) + } + snap := st.Snapshot() + if snap.Status != GoalBlocked { + t.Errorf("status = %q, want blocked", snap.Status) + } + if snap.BlockReason != "missing API key" { + t.Errorf("block reason = %q", snap.BlockReason) + } +} + +func TestGoalBlockedRequiresEvidence(t *testing.T) { + st := NewGoalState() + st.Start("g1", "do the thing", 0) + tool := &GoalBlockedTool{State: st} + + res, err := tool.Execute(context.Background(), "c1", + json.RawMessage(`{"reason":"stuck","evidence":""}`), nil) + if err != nil { + t.Fatalf("Execute: %v", err) + } + if !isErrorResult(res) { + t.Fatal("expected error result when evidence is empty") + } + if snap := st.Snapshot(); snap.Status != GoalActive { + t.Errorf("status = %q, want still active", snap.Status) + } +} + +func TestGoalStateRecordIteration(t *testing.T) { + st := NewGoalState() + st.Start("g1", "obj", 1000) + + st.RecordIteration(100, true) // tool activity resets no-progress + st.RecordIteration(50, false) // no tool activity + st.RecordIteration(50, false) + + snap := st.Snapshot() + if snap.Iterations != 3 { + t.Errorf("iterations = %d, want 3", snap.Iterations) + } + if snap.TokensUsed != 200 { + t.Errorf("tokensUsed = %d, want 200", snap.TokensUsed) + } + if snap.NoProgress != 2 { + t.Errorf("noProgress = %d, want 2", snap.NoProgress) + } + if snap.TokenBudget != 1000 { + t.Errorf("tokenBudget = %d, want 1000", snap.TokenBudget) + } +} + +func TestGoalStateClear(t *testing.T) { + st := NewGoalState() + st.Start("g1", "obj", 0) + st.Clear() + if snap := st.Snapshot(); snap.Status != GoalIdle || snap.Objective != "" { + t.Errorf("after Clear: status=%q objective=%q, want idle/empty", snap.Status, snap.Objective) + } +} + +func TestGoalStateResume(t *testing.T) { + st := NewGoalState() + st.Start("g1", "obj", 1000) + // Simulate a run that tripped a safety guard. + st.RecordIteration(500, false) + st.RecordIteration(500, false) + st.SetStatus(GoalPaused) + + st.Resume() + snap := st.Snapshot() + if snap.Status != GoalActive { + t.Errorf("status = %q, want active after Resume", snap.Status) + } + if snap.Iterations != 0 || snap.NoProgress != 0 || snap.TokensUsed != 0 { + t.Errorf("Resume should clear transient counters: iterations=%d noProgress=%d tokensUsed=%d", + snap.Iterations, snap.NoProgress, snap.TokensUsed) + } + if snap.Objective != "obj" || snap.TokenBudget != 1000 { + t.Errorf("Resume should preserve objective/budget: objective=%q budget=%d", snap.Objective, snap.TokenBudget) + } +} diff --git a/pigo/internal/agenttool/htmlmarkdown.go b/pigo/internal/agenttool/htmlmarkdown.go new file mode 100644 index 0000000..a0473f7 --- /dev/null +++ b/pigo/internal/agenttool/htmlmarkdown.go @@ -0,0 +1,48 @@ +// This file implements the HTML→Markdown reduction used by the webfetch tool +// (US-012, #128). It wraps the JohannesKaufmann/html-to-markdown/v2 library: the +// base plugin already strips head/script/style/link/meta/iframe/noscript/input, +// and we additionally register the remaining page chrome (nav/footer/header/ +// aside/form/svg/template) for removal so only readable content survives. The +// commonmark plugin renders headings, links, lists, code, emphasis, and tables. +package agenttool + +import ( + "sync" + + "github.com/JohannesKaufmann/html-to-markdown/v2/converter" + "github.com/JohannesKaufmann/html-to-markdown/v2/plugin/base" + "github.com/JohannesKaufmann/html-to-markdown/v2/plugin/commonmark" +) + +// chromeElements are dropped whole (element + subtree) in addition to the base +// plugin's defaults: they carry no readable body content for a text extraction. +var chromeElements = []string{ + "nav", "footer", "header", "aside", "svg", "form", "template", +} + +// mdConverter is the shared, configured converter. It is built once — NewConverter +// registers plugins and tag handlers, which is wasteful to repeat per call, and +// the converter is safe for concurrent ConvertString use. +var mdConverter = sync.OnceValue(func() *converter.Converter { + conv := converter.NewConverter( + converter.WithPlugins( + base.NewBasePlugin(), + commonmark.NewCommonmarkPlugin(), + ), + ) + for _, tag := range chromeElements { + conv.Register.TagType(tag, converter.TagTypeRemove, converter.PriorityStandard) + } + return conv +}) + +// htmlToMarkdown converts body (HTML) to a simplified Markdown string. On a +// conversion error (rare — the parser is lenient) it falls back to the raw bytes +// so the caller always gets usable text. +func htmlToMarkdown(body []byte) string { + md, err := mdConverter().ConvertString(string(body)) + if err != nil { + return string(body) + } + return md +} diff --git a/pigo/internal/agenttool/htmlmarkdown_test.go b/pigo/internal/agenttool/htmlmarkdown_test.go new file mode 100644 index 0000000..f4f0f9a --- /dev/null +++ b/pigo/internal/agenttool/htmlmarkdown_test.go @@ -0,0 +1,71 @@ +// Tests for the HTML→Markdown reduction (US-012, #128): headings, links, lists, +// code, and dropped chrome/script elements. The conversion is delegated to the +// html-to-markdown/v2 library; these tests pin the behavior webfetch relies on. +package agenttool + +import ( + "strings" + "testing" +) + +// TestHTMLToMarkdownBasics checks headings, emphasis, and links render. +func TestHTMLToMarkdownBasics(t *testing.T) { + html := ` +

Section

+

Some bold and a link.

+ ` + md := htmlToMarkdown([]byte(html)) + for _, want := range []string{"## Section", "**bold**", "[link](https://go.dev)"} { + if !strings.Contains(md, want) { + t.Errorf("markdown missing %q in:\n%s", want, md) + } + } +} + +// TestHTMLToMarkdownDropsChrome checks script/style/nav/footer content is removed +// while real body content survives. +func TestHTMLToMarkdownDropsChrome(t *testing.T) { + html := ` + + +

real content

+
copyright notice
+ ` + md := htmlToMarkdown([]byte(html)) + if !strings.Contains(md, "real content") { + t.Errorf("body content dropped: %q", md) + } + for _, gone := range []string{"tracker()", ".x{}", "menu links", "copyright notice"} { + if strings.Contains(md, gone) { + t.Errorf("chrome/noise %q leaked into: %q", gone, md) + } + } +} + +// TestHTMLToMarkdownLists checks list items become dashes. +func TestHTMLToMarkdownLists(t *testing.T) { + html := `
  • one
  • two
` + md := htmlToMarkdown([]byte(html)) + if !strings.Contains(md, "one") || !strings.Contains(md, "two") { + t.Errorf("list not rendered: %q", md) + } + if !strings.Contains(md, "- ") { + t.Errorf("list markers missing: %q", md) + } +} + +// TestHTMLToMarkdownInlineSpacing checks spaces around inline elements survive +// (regression: "A link here" must not become "Alinkhere"). +func TestHTMLToMarkdownInlineSpacing(t *testing.T) { + md := htmlToMarkdown([]byte(`

A link here.

`)) + if !strings.Contains(md, "A [link](https://x.io) here.") { + t.Errorf("inline spacing lost: %q", md) + } +} + +// TestHTMLToMarkdownEmptyFallback checks empty input does not panic. +func TestHTMLToMarkdownEmptyFallback(t *testing.T) { + if got := htmlToMarkdown([]byte("")); strings.TrimSpace(got) != "" { + t.Errorf("empty input = %q, want empty", got) + } +} diff --git a/pigo/internal/agenttool/memory_tool.go b/pigo/internal/agenttool/memory_tool.go new file mode 100644 index 0000000..f274393 --- /dev/null +++ b/pigo/internal/agenttool/memory_tool.go @@ -0,0 +1,170 @@ +// This file implements the memory_search AgentTool (issue #477): a read-only +// tool that queries the persistent memory library (internal/memory) by relevance +// using its BM25 full-text index and returns ranked snippets to the model. +// +// Writes are intentionally NOT a separate tool here. Per the SPEC (§4.1/§4.2), +// memory writes reuse the existing Write/Edit file tools, constrained to the +// memory root and carrying the canonical frontmatter (name/description/ +// metadata.type). Traversal protection for those writes lives in the memory +// package (memory.assertSafeComponent) and the write-path plumbing; a dedicated +// memory_write tool would duplicate that. After an off-tool write, the next +// memory_search picks it up automatically because Execute searches with +// ReconcileFirst=true (lazy reconcile). +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/memory" +) + +// memorySearchDefaultLimit is the result cap used when the caller omits limit or +// passes a non-positive value; memorySearchMaxLimit is the hard upper bound. +const ( + memorySearchDefaultLimit = 10 + memorySearchMaxLimit = 50 +) + +// MemorySearchTool searches the persistent memory library by relevance. Store is +// exported so the loop-integration node (#481) can construct the tool with a +// live *memory.Store, mirroring how TodoTool exposes its Store. +type MemorySearchTool struct { + // Store is the persistent memory store. When nil, Execute degrades to a + // friendly no-op result rather than erroring, so a session without memory + // configured still runs. + Store *memory.Store +} + +// memorySearchArgs is the decoded argument shape for memory_search. +type memorySearchArgs struct { + Query string `json:"query"` + Scope string `json:"scope"` + Type string `json:"type"` + Limit int `json:"limit"` +} + +// Name implements AgentTool. +func (t *MemorySearchTool) Name() string { return "memory_search" } + +// Description implements AgentTool. +func (t *MemorySearchTool) Description() string { + return "Search the persistent memory library by relevance (BM25 full-text) " + + "and return ranked snippets from previously saved notes, checkpoints, and " + + "references. Use it to recall context from earlier sessions before " + + "answering or acting. Optional filters: scope (global|projects|sessions|cc), " + + "type (user|feedback|project|reference|checkpoint|progress|notes|free), and " + + "limit (default 10, max 50). Results are ordered most-relevant first." +} + +// Schema implements AgentTool. +func (t *MemorySearchTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Free-text query; tokenized and matched against memory bodies (BM25)."}, + "scope": {"type": "string", "description": "Optional scope filter: global, projects, sessions, or cc."}, + "type": {"type": "string", "description": "Optional type filter: user, feedback, project, reference, checkpoint, progress, notes, or free."}, + "limit": {"type": "integer", "description": "Max results to return (default 10, capped at 50)."} + }, + "required": ["query"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. memory_search is read-only, so it runs in +// the default parallel mode. +func (t *MemorySearchTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} + +// Execute implements AgentTool. It decodes the args, runs a lazily-reconciled +// BM25 search, and formats the ranked hits as text (one line each: +// "[type/scope] path (score) — snippet") with the structured []SearchResult in +// Details. A nil Store or empty query degrades to a friendly no-op result rather +// than a Go error so the loop keeps running. +func (t *MemorySearchTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[memorySearchArgs](args, "memory_search") + if bad != nil { + return *bad, nil + } + + query := strings.TrimSpace(a.Query) + if t.Store == nil { + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent( + "memory_search: no memory store configured; nothing to search.")}, + }, nil + } + if query == "" { + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent( + "memory_search: empty query; provide a search string.")}, + }, nil + } + + limit := a.Limit + if limit <= 0 { + limit = memorySearchDefaultLimit + } + if limit > memorySearchMaxLimit { + limit = memorySearchMaxLimit + } + + results, err := t.Store.Search(query, memory.SearchOptions{ + Scope: strings.TrimSpace(a.Scope), + Type: strings.TrimSpace(a.Type), + Limit: limit, + ReconcileFirst: true, // lazy reconcile so off-tool writes are indexed + // ScoreFloor left at its zero value → package default (0.15). + }) + if err != nil { + return errorResult(fmt.Sprintf("memory_search: %v", err)), nil + } + + if len(results) == 0 { + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent( + fmt.Sprintf("memory_search: no results for %q.", query))}, + Details: results, + }, nil + } + + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(formatMemoryResults(query, results))}, + Details: results, + }, nil +} + +// formatMemoryResults renders ranked hits as a header line plus one line per +// result: "[type/scope] path (score) — snippet". Snippets are whitespace- +// collapsed so a multi-line body stays on a single row. +func formatMemoryResults(query string, results []memory.SearchResult) string { + var b strings.Builder + fmt.Fprintf(&b, "memory_search: %d result(s) for %q, most relevant first:", len(results), query) + for _, r := range results { + typ := string(r.Type) + if typ == "" { + typ = "free" + } + scope := string(r.Scope) + if r.ScopeID != "" { + scope = scope + "/" + r.ScopeID + } + line := fmt.Sprintf("\n[%s/%s] %s (%.3f)", typ, scope, r.Path, r.Score) + if snip := collapseWhitespace(r.Snippet); snip != "" { + line += " — " + snip + } + b.WriteString(line) + } + return b.String() +} + +// collapseWhitespace folds any run of whitespace (including newlines) into a +// single space and trims the ends, keeping a snippet to one line. +func collapseWhitespace(s string) string { + return strings.Join(strings.Fields(s), " ") +} diff --git a/pigo/internal/agenttool/memory_tool_test.go b/pigo/internal/agenttool/memory_tool_test.go new file mode 100644 index 0000000..e43afc2 --- /dev/null +++ b/pigo/internal/agenttool/memory_tool_test.go @@ -0,0 +1,179 @@ +package agenttool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/memory" +) + +// newMemoryStoreWithCorpus opens a *memory.Store over a temp DB + temp mimo root +// and writes a couple of .md files under the layout. It does NOT reconcile — the +// tool's ReconcileFirst=true is expected to index them lazily on first search. +func newMemoryStoreWithCorpus(t *testing.T) *memory.Store { + t.Helper() + base := t.TempDir() + root := filepath.Join(base, "mimo") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatalf("mkdir root: %v", err) + } + writeMemFile(t, root, "permission deadlock encountered during checkpoint save then retry succeeded", + "projects", "proj1", "notes", "rare.md") + writeMemFile(t, root, "unrelated grocery shopping list", + "global", "user", "u1.md") + + dbPath := filepath.Join(base, "sub", "memory.db") + st, err := memory.Open(dbPath, root, "") + if err != nil { + t.Fatalf("memory.Open: %v", err) + } + t.Cleanup(func() { st.Close() }) + return st +} + +func writeMemFile(t *testing.T, root, body string, segs ...string) string { + t.Helper() + full := filepath.Join(append([]string{root}, segs...)...) + if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil { + t.Fatalf("mkdir for %q: %v", full, err) + } + if err := os.WriteFile(full, []byte(body), 0o644); err != nil { + t.Fatalf("write %q: %v", full, err) + } + return filepath.Clean(full) +} + +func runMemorySearch(t *testing.T, tool *MemorySearchTool, args map[string]any) (string, any) { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + res, err := tool.Execute(context.Background(), "call-1", raw, nil) + if err != nil { + t.Fatalf("Execute returned Go error: %v", err) + } + return contentText(res.Content), res.Details +} + +// contentText concatenates the text of every TextContent block in a result. +func contentText(content agentcore.ContentList) string { + var b strings.Builder + for _, c := range content { + if tc, ok := c.(agentcore.TextContent); ok { + b.WriteString(tc.Text) + } + } + return b.String() +} + +func TestMemorySearchToolInterface(t *testing.T) { + tool := &MemorySearchTool{} + if tool.Name() != "memory_search" { + t.Fatalf("Name = %q, want memory_search", tool.Name()) + } + if tool.Description() == "" { + t.Fatal("Description must not be empty") + } + // Schema must be valid JSON declaring query as required. + var schema struct { + Required []string `json:"required"` + } + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Fatalf("Schema is not valid JSON: %v", err) + } + if len(schema.Required) != 1 || schema.Required[0] != "query" { + t.Fatalf("Schema required = %v, want [query]", schema.Required) + } +} + +func TestMemorySearchFindsSnippet(t *testing.T) { + tool := &MemorySearchTool{Store: newMemoryStoreWithCorpus(t)} + + text, details := runMemorySearch(t, tool, map[string]any{"query": "permission deadlock"}) + + if !strings.Contains(text, "rare.md") { + t.Fatalf("expected result text to reference rare.md, got:\n%s", text) + } + if !strings.Contains(strings.ToLower(text), "permission") { + t.Fatalf("expected snippet to mention 'permission', got:\n%s", text) + } + + // Details must carry the structured results (lazy reconcile indexed the file). + results, ok := details.([]memory.SearchResult) + if !ok { + t.Fatalf("Details type = %T, want []memory.SearchResult", details) + } + if len(results) == 0 { + t.Fatal("expected at least one structured result") + } + found := false + for _, r := range results { + if strings.HasSuffix(r.Path, filepath.Join("notes", "rare.md")) { + found = true + } + } + if !found { + t.Fatalf("expected rare.md among structured results, got %+v", results) + } +} + +func TestMemorySearchScopeAndTypeFilter(t *testing.T) { + tool := &MemorySearchTool{Store: newMemoryStoreWithCorpus(t)} + + // Filter to the global/user doc; the projects/notes doc must be excluded even + // though it also matches the shared word. + text, _ := runMemorySearch(t, tool, map[string]any{ + "query": "grocery permission", + "scope": "global", + "type": "user", + }) + if strings.Contains(text, "rare.md") { + t.Fatalf("scope/type filter should exclude rare.md, got:\n%s", text) + } + if !strings.Contains(text, "u1.md") { + t.Fatalf("expected u1.md to match global/user filter, got:\n%s", text) + } +} + +func TestMemorySearchNoResults(t *testing.T) { + tool := &MemorySearchTool{Store: newMemoryStoreWithCorpus(t)} + text, _ := runMemorySearch(t, tool, map[string]any{"query": "zzzznonexistenttoken"}) + if !strings.Contains(text, "no results") { + t.Fatalf("expected a clear empty message, got:\n%s", text) + } +} + +func TestMemorySearchEmptyQueryNoOp(t *testing.T) { + tool := &MemorySearchTool{Store: newMemoryStoreWithCorpus(t)} + text, _ := runMemorySearch(t, tool, map[string]any{"query": " "}) + if !strings.Contains(text, "empty query") { + t.Fatalf("expected empty-query no-op message, got:\n%s", text) + } +} + +func TestMemorySearchNilStoreNoOp(t *testing.T) { + tool := &MemorySearchTool{} // Store nil + text, _ := runMemorySearch(t, tool, map[string]any{"query": "anything"}) + if !strings.Contains(text, "no memory store") { + t.Fatalf("expected nil-store no-op message, got:\n%s", text) + } +} + +func TestMemorySearchInvalidArgs(t *testing.T) { + tool := &MemorySearchTool{Store: newMemoryStoreWithCorpus(t)} + res, err := tool.Execute(context.Background(), "call-1", json.RawMessage(`{"query": 123}`), nil) + if err != nil { + t.Fatalf("Execute returned Go error: %v", err) + } + var text strings.Builder + text.WriteString(contentText(res.Content)) + if !strings.Contains(text.String(), "invalid arguments") { + t.Fatalf("expected invalid-arguments error result, got:\n%s", text.String()) + } +} diff --git a/pigo/internal/agenttool/read_tool.go b/pigo/internal/agenttool/read_tool.go new file mode 100644 index 0000000..e8a5027 --- /dev/null +++ b/pigo/internal/agenttool/read_tool.go @@ -0,0 +1,178 @@ +// This file implements the read tool (US-015): read a file's contents by path, +// with optional line offset/limit, numbered output, and large-file truncation. +// Paths are resolved against a Root and rejected if they escape it (path +// traversal guard) or do not exist. +package agenttool + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// readToolMaxLines caps how many lines a single read returns before truncating +// (protects the model's context from huge files). Callers page with offset. +const readToolMaxLines = 2000 + +// readToolMaxLineLen caps how many bytes of a single line are returned; longer +// lines are truncated with a marker. +const readToolMaxLineLen = 2000 + +// scanBufInit is the initial per-line scanner buffer (it grows on demand up to +// the max). readScanBufMax is generous — a read may page through a file with +// very long lines (minified JS, JSON) that must not error out mid-read. +const ( + scanBufInit = 64 * 1024 + readScanBufMax = 16 * 1024 * 1024 + grepScanBufMax = 1 * 1024 * 1024 +) + +// filePerm / dirPerm are the modes new files and parent directories are created +// with by the write/edit tools (standard non-executable file, traversable dir). +const ( + filePerm os.FileMode = 0o644 + dirPerm os.FileMode = 0o755 +) + +// ReadTool reads text files under Root. It is the first concrete AgentTool. +type ReadTool struct { + // Root is the directory that bounds all reads. A path resolving outside Root + // is rejected. Empty Root defaults to the current working directory. + Root string + // ExtraRoots are additional trusted directories a read may target even though + // they lie outside Root. It exists for the skills directory: pigo advertises + // each skill's absolute SKILL.md path in the system prompt and tells the model + // to read it, so the read tool must permit those paths (they are otherwise + // outside the workspace and rejected). + ExtraRoots []string +} + +// readToolArgs is the decoded argument shape for ReadTool. +type readToolArgs struct { + // Path is the file to read, relative to Root (or absolute within Root). + Path string `json:"path"` + // Offset is the 1-based line to start reading from. 0/1 both mean line 1. + Offset int `json:"offset,omitempty"` + // Limit is the maximum number of lines to return. 0 means the default cap. + Limit int `json:"limit,omitempty"` +} + +// Name implements AgentTool. +func (t *ReadTool) Name() string { return "read" } + +// Description implements AgentTool. +func (t *ReadTool) Description() string { + return "Read a text file's contents by path, with optional line offset/limit. " + + "Output is line-numbered; very large files are truncated." +} + +// Schema implements AgentTool. +func (t *ReadTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "path": {"type": "string", "description": "File path to read, relative to the workspace root."}, + "offset": {"type": "integer", "description": "1-based line number to start from.", "minimum": 0}, + "limit": {"type": "integer", "description": "Maximum number of lines to return.", "minimum": 0} + }, + "required": ["path"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Reads are side-effect free → parallel. +func (t *ReadTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} + +// resolvePath resolves p against Root (or any ExtraRoots) via the shared +// resolveWithin boundary policy, so every file tool enforces the same +// workspace-escape guard while the read tool can also reach trusted extra roots. +func (t *ReadTool) resolvePath(p string) (string, error) { + if len(t.ExtraRoots) == 0 { + return resolveWithin(t.Root, p) + } + return resolveWithinAny(append([]string{t.Root}, t.ExtraRoots...), p) +} + +// Execute implements AgentTool. It never returns a Go error for a read failure +// (bad path, missing file); those are encoded as error results so the model can +// react. The returned error is reserved for argument decode failures. +func (t *ReadTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[readToolArgs](args, "read") + if bad != nil { + return *bad, nil + } + if a.Path == "" { + return errorResult("read: path is required"), nil + } + full, err := t.resolvePath(a.Path) + if err != nil { + return errorResult("read: " + err.Error()), nil + } + info, err := os.Stat(full) + if err != nil { + if os.IsNotExist(err) { + return errorResult(fmt.Sprintf("read: file %q does not exist", a.Path)), nil + } + return errorResult(fmt.Sprintf("read: cannot stat %q: %v", a.Path, err)), nil + } + if info.IsDir() { + return errorResult(fmt.Sprintf("read: %q is a directory, not a file", a.Path)), nil + } + + f, err := os.Open(full) + if err != nil { + return errorResult(fmt.Sprintf("read: cannot open %q: %v", a.Path, err)), nil + } + defer f.Close() + + text, truncated := readNumbered(f, a.Offset, a.Limit) + if truncated { + text += fmt.Sprintf("\n... (output truncated at %d lines; use offset to read more)", readToolMaxLines) + } + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(text)}}, nil +} + +// readNumbered reads lines from r starting at 1-based offset, returning at most +// limit lines (capped at readToolMaxLines), each prefixed with its line number. +// The bool reports whether the output was truncated by the cap. +func readNumbered(r io.Reader, offset, limit int) (string, bool) { + if offset < 1 { + offset = 1 + } + max := limit + if max <= 0 || max > readToolMaxLines { + max = readToolMaxLines + } + + sc := bufio.NewScanner(r) + sc.Buffer(make([]byte, 0, scanBufInit), readScanBufMax) + var b strings.Builder + lineNo := 0 + emitted := 0 + truncated := false + for sc.Scan() { + lineNo++ + if lineNo < offset { + continue + } + if emitted >= max { + truncated = true + break + } + line := sc.Text() + if len(line) > readToolMaxLineLen { + line = line[:readToolMaxLineLen] + "… (line truncated)" + } + fmt.Fprintf(&b, "%6d\t%s\n", lineNo, line) + emitted++ + } + return b.String(), truncated +} diff --git a/pigo/internal/agenttool/read_tool_test.go b/pigo/internal/agenttool/read_tool_test.go new file mode 100644 index 0000000..08eb237 --- /dev/null +++ b/pigo/internal/agenttool/read_tool_test.go @@ -0,0 +1,204 @@ +package agenttool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func runRead(t *testing.T, tool *ReadTool, args map[string]any) (agentcore.AgentToolResult, bool) { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + res, gerr := tool.Execute(context.Background(), "call-1", raw, nil) + if gerr != nil { + t.Fatalf("execute returned go error: %v", gerr) + } + return res, false +} + +func resultText(res agentcore.AgentToolResult) string { + var b strings.Builder + for _, c := range res.Content { + if tc, ok := c.(agentcore.TextContent); ok { + b.WriteString(tc.Text) + } + } + return b.String() +} + +func TestReadToolBasic(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "hello.txt") + if err := os.WriteFile(path, []byte("line one\nline two\nline three\n"), 0o644); err != nil { + t.Fatalf("write: %v", err) + } + tool := &ReadTool{Root: dir} + res, _ := runRead(t, tool, map[string]any{"path": "hello.txt"}) + text := resultText(res) + if !strings.Contains(text, "line one") || !strings.Contains(text, "line three") { + t.Errorf("missing content: %q", text) + } + // Line numbers present. + if !strings.Contains(text, "1\tline one") || !strings.Contains(text, "3\tline three") { + t.Errorf("missing line numbers: %q", text) + } +} + +func TestReadToolOffsetLimit(t *testing.T) { + dir := t.TempDir() + var sb strings.Builder + for i := 1; i <= 10; i++ { + sb.WriteString("row\n") + } + path := filepath.Join(dir, "rows.txt") + if err := os.WriteFile(path, []byte(sb.String()), 0o644); err != nil { + t.Fatalf("write: %v", err) + } + tool := &ReadTool{Root: dir} + res, _ := runRead(t, tool, map[string]any{"path": "rows.txt", "offset": 3, "limit": 2}) + text := resultText(res) + // Should include line numbers 3 and 4, not 1,2,5. + if !strings.Contains(text, "3\trow") || !strings.Contains(text, "4\trow") { + t.Errorf("offset/limit window wrong: %q", text) + } + if strings.Contains(text, "2\trow") || strings.Contains(text, "5\trow") { + t.Errorf("offset/limit leaked outside window: %q", text) + } +} + +func TestReadToolMissingFile(t *testing.T) { + tool := &ReadTool{Root: t.TempDir()} + res, _ := runRead(t, tool, map[string]any{"path": "nope.txt"}) + if !strings.Contains(resultText(res), "does not exist") { + t.Errorf("expected does-not-exist error, got %q", resultText(res)) + } +} + +func TestReadToolPathTraversal(t *testing.T) { + dir := t.TempDir() + // A secret sits outside the root. + parent := filepath.Dir(dir) + secret := filepath.Join(parent, "secret.txt") + if err := os.WriteFile(secret, []byte("top secret"), 0o644); err != nil { + t.Fatalf("write secret: %v", err) + } + defer os.Remove(secret) + + tool := &ReadTool{Root: dir} + res, _ := runRead(t, tool, map[string]any{"path": "../secret.txt"}) + text := resultText(res) + if strings.Contains(text, "top secret") { + t.Fatal("path traversal escaped the root!") + } + if !strings.Contains(text, "outside the workspace root") { + t.Errorf("expected boundary error, got %q", text) + } +} + +func TestReadToolDirectory(t *testing.T) { + dir := t.TempDir() + sub := filepath.Join(dir, "subdir") + if err := os.Mkdir(sub, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + tool := &ReadTool{Root: dir} + res, _ := runRead(t, tool, map[string]any{"path": "subdir"}) + if !strings.Contains(resultText(res), "is a directory") { + t.Errorf("expected directory error, got %q", resultText(res)) + } +} + +func TestReadToolTruncation(t *testing.T) { + dir := t.TempDir() + var sb strings.Builder + for i := 0; i < readToolMaxLines+50; i++ { + sb.WriteString("x\n") + } + path := filepath.Join(dir, "big.txt") + if err := os.WriteFile(path, []byte(sb.String()), 0o644); err != nil { + t.Fatalf("write: %v", err) + } + tool := &ReadTool{Root: dir} + res, _ := runRead(t, tool, map[string]any{"path": "big.txt"}) + if !strings.Contains(resultText(res), "output truncated") { + t.Error("expected truncation notice for oversized file") + } +} + +func TestReadToolMissingPathArg(t *testing.T) { + tool := &ReadTool{Root: t.TempDir()} + res, _ := runRead(t, tool, map[string]any{}) + if !strings.Contains(resultText(res), "path is required") { + t.Errorf("expected path-required error, got %q", resultText(res)) + } +} + +func TestReadToolExtraRootsAllowsTrustedOutsidePath(t *testing.T) { + work := t.TempDir() + // A skill file lives OUTSIDE the workspace root (mirrors ~/.agents/skills). + skills := t.TempDir() + skillFile := filepath.Join(skills, "weather", "SKILL.md") + if err := os.MkdirAll(filepath.Dir(skillFile), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + if err := os.WriteFile(skillFile, []byte("skill body"), 0o644); err != nil { + t.Fatalf("write skill: %v", err) + } + + // Without ExtraRoots the absolute skill path is rejected as out-of-workspace. + bounded := &ReadTool{Root: work} + res, _ := runRead(t, bounded, map[string]any{"path": skillFile}) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Fatalf("expected boundary rejection without ExtraRoots, got %q", resultText(res)) + } + + // With the skills dir as an extra root the same read succeeds. + tool := &ReadTool{Root: work, ExtraRoots: []string{skills}} + res, _ = runRead(t, tool, map[string]any{"path": skillFile}) + if !strings.Contains(resultText(res), "skill body") { + t.Fatalf("expected skill contents with ExtraRoots, got %q", resultText(res)) + } +} + +func TestReadToolExtraRootsStillBlocksUntrustedPath(t *testing.T) { + work := t.TempDir() + skills := t.TempDir() + // A secret sits outside BOTH the workspace root and the extra root. + other := t.TempDir() + secret := filepath.Join(other, "secret.txt") + if err := os.WriteFile(secret, []byte("top secret"), 0o644); err != nil { + t.Fatalf("write secret: %v", err) + } + + tool := &ReadTool{Root: work, ExtraRoots: []string{skills}} + res, _ := runRead(t, tool, map[string]any{"path": secret}) + text := resultText(res) + if strings.Contains(text, "top secret") { + t.Fatal("read escaped both roots!") + } + if !strings.Contains(text, "outside the workspace root") { + t.Errorf("expected boundary error for untrusted path, got %q", text) + } +} + +func TestReadToolSchemaAndMode(t *testing.T) { + tool := &ReadTool{} + if tool.Name() != "read" { + t.Errorf("name = %q", tool.Name()) + } + if tool.ExecutionMode() != agentcore.ToolExecutionParallel { + t.Errorf("read should be parallel") + } + var schema map[string]any + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Errorf("schema not valid JSON: %v", err) + } +} diff --git a/pigo/internal/agenttool/registry.go b/pigo/internal/agenttool/registry.go new file mode 100644 index 0000000..b438f8d --- /dev/null +++ b/pigo/internal/agenttool/registry.go @@ -0,0 +1,220 @@ +// This file implements the tool registry (US-014): tools are registered by +// name, and their arguments are validated against a per-tool JSON Schema +// (santhosh-tekuri/jsonschema v6) before execution. Validation failures are +// turned into a field-level error tool result rather than a Go error, so the +// model receives actionable feedback in the loop. +package agenttool + +import ( + "bytes" + "encoding/json" + "fmt" + "sort" + "strings" + "sync" + + "github.com/santhosh-tekuri/jsonschema/v6" + "github.com/smallnest/pigo/internal/agentcore" + "golang.org/x/text/language" + "golang.org/x/text/message" +) + +// schemaPrinter renders jsonschema error kinds. LocalizedString dereferences +// the printer, so it must be non-nil. +var schemaPrinter = message.NewPrinter(language.English) + +// ToolRegistry stores tools by name and validates call arguments against each +// tool's declared JSON Schema. It is safe for concurrent use. +type ToolRegistry struct { + mu sync.RWMutex + tools map[string]agentcore.AgentTool + compiled map[string]*jsonschema.Schema +} + +// NewToolRegistry returns an empty registry. +func NewToolRegistry() *ToolRegistry { + return &ToolRegistry{ + tools: make(map[string]agentcore.AgentTool), + compiled: make(map[string]*jsonschema.Schema), + } +} + +// Register adds a tool, compiling its JSON Schema up front so bad schemas fail +// at registration rather than on first call. A duplicate name is an error. A +// tool whose Schema() is empty is registered with no validation. +func (r *ToolRegistry) Register(tool agentcore.AgentTool) error { + name := tool.Name() + if name == "" { + return fmt.Errorf("registry: tool has empty name") + } + r.mu.Lock() + defer r.mu.Unlock() + if _, exists := r.tools[name]; exists { + return fmt.Errorf("registry: tool %q already registered", name) + } + if raw := tool.Schema(); len(bytes.TrimSpace(raw)) > 0 && !bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + sch, err := compileSchema(name, raw) + if err != nil { + return fmt.Errorf("registry: tool %q schema: %w", name, err) + } + r.compiled[name] = sch + } + r.tools[name] = tool + return nil +} + +// Get returns the tool registered under name and whether it was found. +func (r *ToolRegistry) Get(name string) (agentcore.AgentTool, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + t, ok := r.tools[name] + return t, ok +} + +// List returns all registered tools sorted by name (stable ordering for +// deterministic provider tool declarations). +func (r *ToolRegistry) List() []agentcore.AgentTool { + r.mu.RLock() + defer r.mu.RUnlock() + out := make([]agentcore.AgentTool, 0, len(r.tools)) + for _, t := range r.tools { + out = append(out, t) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name() < out[j].Name() }) + return out +} + +// FieldError is a single validation failure located at a JSON-pointer path +// within the arguments. +type FieldError struct { + Field string `json:"field"` // JSON pointer, e.g. "/path" or "" for root + Message string `json:"message"` // human-readable reason +} + +// Validate checks args against the tool's compiled schema. It returns nil when +// the tool has no schema or the arguments are valid; otherwise it returns the +// flattened field-level errors. An unknown tool name is reported as a single +// root-level error. +func (r *ToolRegistry) Validate(name string, args json.RawMessage) []FieldError { + r.mu.RLock() + _, known := r.tools[name] + sch, hasSchema := r.compiled[name] + r.mu.RUnlock() + + if !known { + return []FieldError{{Field: "", Message: fmt.Sprintf("unknown tool %q", name)}} + } + if !hasSchema { + return nil + } + + var inst any + dec := json.NewDecoder(bytes.NewReader(nonEmptyJSON(args))) + dec.UseNumber() + if err := dec.Decode(&inst); err != nil { + return []FieldError{{Field: "", Message: fmt.Sprintf("arguments are not valid JSON: %v", err)}} + } + + if err := sch.Validate(inst); err != nil { + var verr *jsonschema.ValidationError + if as := asValidationError(err); as != nil { + verr = as + } + if verr != nil { + return flattenValidationError(verr) + } + return []FieldError{{Field: "", Message: err.Error()}} + } + return nil +} + +// ValidationErrorResult builds an error AgentToolResult describing the given +// field errors, for the loop to hand back to the model (FR: field-level error +// tool result). Terminate is left nil (a validation failure never ends the run). +func ValidationErrorResult(toolName string, errs []FieldError) agentcore.AgentToolResult { + var b strings.Builder + fmt.Fprintf(&b, "Invalid arguments for tool %q:\n", toolName) + for _, e := range errs { + field := e.Field + if field == "" { + field = "(root)" + } + fmt.Fprintf(&b, " - %s: %s\n", field, e.Message) + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(strings.TrimRight(b.String(), "\n"))}, + Details: errs, + } +} + +// compileSchema compiles a raw JSON Schema document held in memory. +func compileSchema(name string, raw json.RawMessage) (*jsonschema.Schema, error) { + doc, err := jsonschema.UnmarshalJSON(bytes.NewReader(raw)) + if err != nil { + return nil, err + } + c := jsonschema.NewCompiler() + // A synthetic in-memory URL; each tool gets its own so schemas never clash. + loc := "mem:///" + name + ".json" + if err := c.AddResource(loc, doc); err != nil { + return nil, err + } + return c.Compile(loc) +} + +// flattenValidationError walks the ValidationError tree and returns one +// FieldError per leaf cause (the most specific failures), falling back to the +// node itself when it has no causes. +func flattenValidationError(e *jsonschema.ValidationError) []FieldError { + var out []FieldError + var walk func(n *jsonschema.ValidationError) + walk = func(n *jsonschema.ValidationError) { + if len(n.Causes) == 0 { + out = append(out, FieldError{ + Field: jsonPointer(n.InstanceLocation), + Message: n.ErrorKind.LocalizedString(schemaPrinter), + }) + return + } + for _, c := range n.Causes { + walk(c) + } + } + walk(e) + if len(out) == 0 { + out = append(out, FieldError{Field: jsonPointer(e.InstanceLocation), Message: e.Error()}) + } + return out +} + +// jsonPointer renders an instance-location path as a JSON pointer. +func jsonPointer(loc []string) string { + if len(loc) == 0 { + return "" + } + var b strings.Builder + for _, tok := range loc { + b.WriteByte('/') + tok = strings.ReplaceAll(tok, "~", "~0") + tok = strings.ReplaceAll(tok, "/", "~1") + b.WriteString(tok) + } + return b.String() +} + +// asValidationError extracts a *jsonschema.ValidationError from err if present. +func asValidationError(err error) *jsonschema.ValidationError { + if verr, ok := err.(*jsonschema.ValidationError); ok { + return verr + } + return nil +} + +// nonEmptyJSON treats empty arguments as an empty object so schemas with only +// optional properties validate, and "required" violations are reported. +func nonEmptyJSON(args json.RawMessage) []byte { + if len(bytes.TrimSpace(args)) == 0 { + return []byte("{}") + } + return args +} diff --git a/pigo/internal/agenttool/registry_test.go b/pigo/internal/agenttool/registry_test.go new file mode 100644 index 0000000..2dc236e --- /dev/null +++ b/pigo/internal/agenttool/registry_test.go @@ -0,0 +1,138 @@ +package agenttool + +import ( + "context" + "encoding/json" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// stubTool is a minimal AgentTool for registry tests. +type stubTool struct { + name string + schema string +} + +func (s stubTool) Name() string { return s.name } +func (s stubTool) Description() string { return "stub" } +func (s stubTool) Schema() json.RawMessage { return json.RawMessage(s.schema) } +func (s stubTool) ExecutionMode() agentcore.ToolExecutionMode { return agentcore.ToolExecutionParallel } +func (s stubTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, nil +} + +const personSchema = `{ + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"} + }, + "required": ["name"], + "additionalProperties": false +}` + +func newTestRegistry(t *testing.T) *ToolRegistry { + t.Helper() + r := NewToolRegistry() + if err := r.Register(stubTool{name: "person", schema: personSchema}); err != nil { + t.Fatalf("register: %v", err) + } + return r +} + +func TestRegistryRegisterAndGet(t *testing.T) { + r := newTestRegistry(t) + tool, ok := r.Get("person") + if !ok || tool.Name() != "person" { + t.Fatalf("Get(person) failed: %v %v", tool, ok) + } + if _, ok := r.Get("missing"); ok { + t.Error("Get(missing) should report not found") + } + if got := r.List(); len(got) != 1 || got[0].Name() != "person" { + t.Errorf("List wrong: %+v", got) + } +} + +func TestRegistryDuplicateRejected(t *testing.T) { + r := newTestRegistry(t) + if err := r.Register(stubTool{name: "person", schema: personSchema}); err == nil { + t.Fatal("expected duplicate registration to error") + } +} + +func TestRegistryValidArgs(t *testing.T) { + r := newTestRegistry(t) + errs := r.Validate("person", json.RawMessage(`{"name":"ada","age":36}`)) + if errs != nil { + t.Fatalf("valid args reported errors: %+v", errs) + } +} + +func TestRegistryMissingRequiredField(t *testing.T) { + r := newTestRegistry(t) + errs := r.Validate("person", json.RawMessage(`{"age":36}`)) + if len(errs) == 0 { + t.Fatal("expected error for missing required field 'name'") + } +} + +func TestRegistryTypeError(t *testing.T) { + r := newTestRegistry(t) + errs := r.Validate("person", json.RawMessage(`{"name":"ada","age":"old"}`)) + if len(errs) == 0 { + t.Fatal("expected type error for age") + } + // The offending field should be located at /age. + found := false + for _, e := range errs { + if e.Field == "/age" { + found = true + } + } + if !found { + t.Errorf("expected a field error at /age, got %+v", errs) + } +} + +func TestRegistryUnknownTool(t *testing.T) { + r := newTestRegistry(t) + errs := r.Validate("nope", json.RawMessage(`{}`)) + if len(errs) != 1 || errs[0].Field != "" { + t.Fatalf("expected single root error for unknown tool, got %+v", errs) + } +} + +func TestRegistryNoSchemaSkipsValidation(t *testing.T) { + r := NewToolRegistry() + if err := r.Register(stubTool{name: "free", schema: ""}); err != nil { + t.Fatalf("register: %v", err) + } + if errs := r.Validate("free", json.RawMessage(`{"anything":true}`)); errs != nil { + t.Fatalf("no-schema tool should skip validation, got %+v", errs) + } +} + +func TestValidationErrorResultShape(t *testing.T) { + r := newTestRegistry(t) + errs := r.Validate("person", json.RawMessage(`{"age":36}`)) + res := ValidationErrorResult("person", errs) + if len(res.Content) == 0 { + t.Fatal("expected content in validation error result") + } + if _, ok := res.Details.([]FieldError); !ok { + t.Errorf("expected Details to carry []FieldError, got %T", res.Details) + } + if res.Terminate != nil { + t.Error("validation failure must not terminate the run") + } +} + +func TestRegistryBadSchemaRejectedAtRegister(t *testing.T) { + r := NewToolRegistry() + err := r.Register(stubTool{name: "bad", schema: `{"type": 123}`}) + if err == nil { + t.Fatal("expected invalid schema to fail at registration") + } +} diff --git a/pigo/internal/agenttool/search_tool.go b/pigo/internal/agenttool/search_tool.go new file mode 100644 index 0000000..68a8d13 --- /dev/null +++ b/pigo/internal/agenttool/search_tool.go @@ -0,0 +1,492 @@ +// This file implements the search tools (US-019): grep (search file contents by +// regexp with optional glob filtering), find (locate files by name glob), and ls +// (list a directory, distinguishing files from directories). All three resolve +// paths against a Root with the same boundary guard as the other tools and skip +// paths ignored by the workspace .gitignore. They are read-only → parallel. +package agenttool + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// searchMaxResults caps the number of matches/entries any single search returns +// so a broad query cannot flood the model's context. +const searchMaxResults = 1000 + +// resolveWithin resolves p against root and verifies it stays within it. It is +// the single workspace-boundary policy shared by every file tool: the search +// tools call it directly, and ReadTool/WriteTool/EditTool.resolvePath delegate +// to it, so the path-traversal guard lives in exactly one place. +func resolveWithin(root, p string) (string, error) { + if root == "" { + wd, err := os.Getwd() + if err != nil { + return "", fmt.Errorf("cannot determine working directory: %w", err) + } + root = wd + } + absRoot, err := filepath.Abs(root) + if err != nil { + return "", fmt.Errorf("invalid root: %w", err) + } + var full string + if filepath.IsAbs(p) { + full = filepath.Clean(p) + } else { + full = filepath.Join(absRoot, p) + } + rel, err := filepath.Rel(absRoot, full) + if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("path %q is outside the workspace root", p) + } + return full, nil +} + +// resolveWithinAny resolves p against the first root that contains it, trying +// roots in order. It exists so the file tools can additionally permit trusted +// out-of-workspace roots (the skills directory) whose absolute SKILL.md paths +// pigo itself advertises in the system prompt: without this, the workspace guard +// would reject the very paths the model is instructed to read or author. Empty +// roots are skipped; if none contain p, the standard workspace-escape error is +// returned. +func resolveWithinAny(roots []string, p string) (string, error) { + var lastErr error + for _, root := range roots { + if root == "" { + continue + } + full, err := resolveWithin(root, p) + if err == nil { + return full, nil + } + lastErr = err + } + if lastErr != nil { + return "", lastErr + } + // No usable roots supplied: fall back to the default (cwd) policy. + return resolveWithin("", p) +} + +// gitignore is a minimal .gitignore matcher. It supports the common subset: +// blank lines and #-comments are skipped; a leading "/" anchors to the root; +// a trailing "/" matches directories only; "!" negation re-includes; and plain +// patterns match by base name or path via filepath.Match. It is intentionally +// not a full gitignore implementation (no "**" spanning, no nested .gitignore). +type gitignore struct { + rules []ignoreRule + // hasSegmentRule is true when at least one rule matches by path segment + // (non-anchored, no "/"). Only then does ignored() need to split relPath into + // segments, so the common all-anchored case skips the split entirely. + hasSegmentRule bool +} + +type ignoreRule struct { + pattern string + negate bool + dirOnly bool + anchored bool + // matchFull is precomputed at load time: an anchored pattern, or one that + // contains a "/", matches against the full relative path; otherwise the rule + // matches by base name or any single path segment. Hoisting this out of the + // per-file loop avoids a strings.Contains scan for every file × rule. + matchFull bool +} + +// loadGitignore reads root/.gitignore. A missing file yields an empty matcher +// (matches nothing), never an error. +func loadGitignore(root string) *gitignore { + gi := &gitignore{} + data, err := os.ReadFile(filepath.Join(root, ".gitignore")) + if err != nil { + return gi + } + sc := bufio.NewScanner(strings.NewReader(string(data))) + for sc.Scan() { + line := strings.TrimRight(sc.Text(), " ") + if line == "" || strings.HasPrefix(line, "#") { + continue + } + r := ignoreRule{} + if strings.HasPrefix(line, "!") { + r.negate = true + line = line[1:] + } + if strings.HasSuffix(line, "/") { + r.dirOnly = true + line = strings.TrimSuffix(line, "/") + } + if strings.HasPrefix(line, "/") { + r.anchored = true + line = strings.TrimPrefix(line, "/") + } + if line == "" { + continue + } + r.pattern = line + r.matchFull = r.anchored || strings.Contains(line, "/") + if !r.matchFull { + gi.hasSegmentRule = true + } + gi.rules = append(gi.rules, r) + } + return gi +} + +// ignored reports whether relPath (slash-separated, relative to root) is ignored. +// isDir refines dir-only rules. Later rules win, so a negation can re-include. +// +// The relPath is split into segments at most once per call (only when a +// segment-matching rule exists), rather than re-splitting inside the rule loop: +// this keeps the per-file cost O(rules) instead of O(rules × pathSegments), +// which matters because ignored() is called for every entry of a WalkDir. +func (g *gitignore) ignored(relPath string, isDir bool) bool { + relPath = filepath.ToSlash(relPath) + base := relPath + if i := strings.LastIndex(relPath, "/"); i >= 0 { + base = relPath[i+1:] + } + var segs []string + if g.hasSegmentRule { + segs = strings.Split(relPath, "/") + } + result := false + for _, r := range g.rules { + if r.dirOnly && !isDir { + continue + } + var match bool + if r.matchFull { + match, _ = filepath.Match(r.pattern, relPath) + } else { + match, _ = filepath.Match(r.pattern, base) + if !match { + // A non-anchored pattern also matches any path component, + // so an ignored directory hides everything beneath it. + for _, seg := range segs { + if ok, _ := filepath.Match(r.pattern, seg); ok { + match = true + break + } + } + } + } + if match { + result = !r.negate + } + } + return result +} + +// GrepTool searches file contents by regexp under Root, honoring .gitignore. +type GrepTool struct { + // Root bounds the search; empty defaults to the current working directory. + Root string +} + +type grepToolArgs struct { + // Pattern is the regexp to search for (Go regexp syntax). + Pattern string `json:"pattern"` + // Path optionally scopes the search to a subdirectory (relative to Root). + Path string `json:"path,omitempty"` + // Glob optionally filters files by base-name glob (e.g. "*.go"). + Glob string `json:"glob,omitempty"` +} + +func (t *GrepTool) Name() string { return "grep" } +func (t *GrepTool) Description() string { + return "Search file contents by regular expression under the workspace, " + + "optionally filtering files by glob. Skips .gitignore'd paths." +} +func (t *GrepTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} +func (t *GrepTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "pattern": {"type": "string", "description": "Regular expression to search for."}, + "path": {"type": "string", "description": "Subdirectory to scope the search to (relative to the workspace root)."}, + "glob": {"type": "string", "description": "Filter files by base-name glob, e.g. *.go."} + }, + "required": ["pattern"], + "additionalProperties": false +}`) +} + +func (t *GrepTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[grepToolArgs](args, "grep") + if bad != nil { + return *bad, nil + } + if a.Pattern == "" { + return errorResult("grep: pattern is required"), nil + } + re, err := regexp.Compile(a.Pattern) + if err != nil { + return errorResult(fmt.Sprintf("grep: invalid pattern: %v", err)), nil + } + root, err := resolveWithin(t.Root, "") + if err != nil { + return errorResult("grep: " + err.Error()), nil + } + start := root + if a.Path != "" { + if start, err = resolveWithin(t.Root, a.Path); err != nil { + return errorResult("grep: " + err.Error()), nil + } + } + gi := loadGitignore(root) + + var matches []string + count := 0 + walkErr := filepath.WalkDir(start, func(path string, d os.DirEntry, err error) error { + if err != nil { + return nil // skip unreadable entries + } + rel, _ := filepath.Rel(root, path) + if rel == "." { + return nil + } + if rel == ".git" || strings.HasPrefix(rel, ".git"+string(filepath.Separator)) { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if gi.ignored(rel, d.IsDir()) { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if d.IsDir() { + return nil + } + if a.Glob != "" { + if ok, _ := filepath.Match(a.Glob, d.Name()); !ok { + return nil + } + } + f, err := os.Open(path) + if err != nil { + return nil + } + defer f.Close() + sc := bufio.NewScanner(f) + sc.Buffer(make([]byte, 0, scanBufInit), grepScanBufMax) + lineNo := 0 + for sc.Scan() { + lineNo++ + line := sc.Text() + if re.MatchString(line) { + matches = append(matches, fmt.Sprintf("%s:%d:%s", rel, lineNo, line)) + count++ + if count >= searchMaxResults { + return filepath.SkipAll + } + } + } + return nil + }) + if walkErr != nil { + return errorResult(fmt.Sprintf("grep: %v", walkErr)), nil + } + + msg := fmt.Sprintf("%d match(es) for %q", len(matches), a.Pattern) + if len(matches) > 0 { + msg += "\n" + strings.Join(matches, "\n") + } + if count >= searchMaxResults { + msg += fmt.Sprintf("\n[truncated at %d matches]", searchMaxResults) + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + Details: map[string]any{"matches": len(matches)}, + }, nil +} + +// FindTool locates files by base-name glob under Root, honoring .gitignore. +type FindTool struct { + // Root bounds the search; empty defaults to the current working directory. + Root string +} + +type findToolArgs struct { + // Glob is the base-name glob to match (e.g. "*.go"). + Glob string `json:"glob"` + // Path optionally scopes the search to a subdirectory (relative to Root). + Path string `json:"path,omitempty"` +} + +func (t *FindTool) Name() string { return "find" } +func (t *FindTool) Description() string { + return "Find files by base-name glob under the workspace. Skips .gitignore'd paths." +} +func (t *FindTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} +func (t *FindTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "glob": {"type": "string", "description": "Base-name glob to match, e.g. *.go."}, + "path": {"type": "string", "description": "Subdirectory to scope the search to (relative to the workspace root)."} + }, + "required": ["glob"], + "additionalProperties": false +}`) +} + +func (t *FindTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[findToolArgs](args, "find") + if bad != nil { + return *bad, nil + } + if a.Glob == "" { + return errorResult("find: glob is required"), nil + } + root, err := resolveWithin(t.Root, "") + if err != nil { + return errorResult("find: " + err.Error()), nil + } + start := root + if a.Path != "" { + if start, err = resolveWithin(t.Root, a.Path); err != nil { + return errorResult("find: " + err.Error()), nil + } + } + gi := loadGitignore(root) + + var found []string + walkErr := filepath.WalkDir(start, func(path string, d os.DirEntry, err error) error { + if err != nil { + return nil + } + rel, _ := filepath.Rel(root, path) + if rel == "." { + return nil + } + if rel == ".git" || strings.HasPrefix(rel, ".git"+string(filepath.Separator)) { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if gi.ignored(rel, d.IsDir()) { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + if d.IsDir() { + return nil + } + if ok, _ := filepath.Match(a.Glob, d.Name()); ok { + found = append(found, rel) + if len(found) >= searchMaxResults { + return filepath.SkipAll + } + } + return nil + }) + if walkErr != nil { + return errorResult(fmt.Sprintf("find: %v", walkErr)), nil + } + sort.Strings(found) + + msg := fmt.Sprintf("%d file(s) matching %q", len(found), a.Glob) + if len(found) > 0 { + msg += "\n" + strings.Join(found, "\n") + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + Details: map[string]any{"count": len(found)}, + }, nil +} + +// LsTool lists the entries of a directory, distinguishing files from directories. +type LsTool struct { + // Root bounds the listing; empty defaults to the current working directory. + Root string +} + +type lsToolArgs struct { + // Path is the directory to list, relative to Root (empty = Root itself). + Path string `json:"path,omitempty"` +} + +func (t *LsTool) Name() string { return "ls" } +func (t *LsTool) Description() string { + return "List a directory's entries, marking directories with a trailing slash." +} +func (t *LsTool) ExecutionMode() agentcore.ToolExecutionMode { return agentcore.ToolExecutionParallel } +func (t *LsTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "path": {"type": "string", "description": "Directory to list, relative to the workspace root."} + }, + "additionalProperties": false +}`) +} + +func (t *LsTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[lsToolArgs](args, "ls") + if bad != nil { + return *bad, nil + } + full, err := resolveWithin(t.Root, a.Path) + if err != nil { + return errorResult("ls: " + err.Error()), nil + } + info, err := os.Stat(full) + if err != nil { + if os.IsNotExist(err) { + return errorResult(fmt.Sprintf("ls: %q does not exist", a.Path)), nil + } + return errorResult(fmt.Sprintf("ls: cannot stat %q: %v", a.Path, err)), nil + } + if !info.IsDir() { + return errorResult(fmt.Sprintf("ls: %q is not a directory", a.Path)), nil + } + entries, err := os.ReadDir(full) + if err != nil { + return errorResult(fmt.Sprintf("ls: cannot read %q: %v", a.Path, err)), nil + } + + var dirs, files []string + for _, e := range entries { + if e.IsDir() { + dirs = append(dirs, e.Name()+"/") + } else { + files = append(files, e.Name()) + } + } + sort.Strings(dirs) + sort.Strings(files) + lines := append(dirs, files...) + + label := a.Path + if label == "" { + label = "." + } + msg := fmt.Sprintf("%s (%d dir(s), %d file(s))", label, len(dirs), len(files)) + if len(lines) > 0 { + msg += "\n" + strings.Join(lines, "\n") + } + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + Details: map[string]any{"dirs": len(dirs), "files": len(files)}, + }, nil +} diff --git a/pigo/internal/agenttool/search_tool_test.go b/pigo/internal/agenttool/search_tool_test.go new file mode 100644 index 0000000..00330dc --- /dev/null +++ b/pigo/internal/agenttool/search_tool_test.go @@ -0,0 +1,225 @@ +package agenttool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func runSearch(t *testing.T, tool agentcore.AgentTool, args map[string]any) agentcore.AgentToolResult { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + res, gerr := tool.Execute(context.Background(), "call-1", raw, nil) + if gerr != nil { + t.Fatalf("execute returned go error: %v", gerr) + } + return res +} + +// seedTree writes a small directory tree with a .gitignore for the search tests. +func seedTree(t *testing.T) string { + t.Helper() + dir := t.TempDir() + mustWrite := func(rel, content string) { + p := filepath.Join(dir, rel) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatalf("mkdir %s: %v", rel, err) + } + if err := os.WriteFile(p, []byte(content), 0o644); err != nil { + t.Fatalf("write %s: %v", rel, err) + } + } + mustWrite("main.go", "package main\nfunc main() { hello() }\n") + mustWrite("util.go", "package main\nfunc hello() {}\n") + mustWrite("README.md", "# project\nhello world\n") + mustWrite("sub/deep.go", "package sub\n// hello from sub\n") + mustWrite("build/generated.go", "package build\nfunc hello() {}\n") + mustWrite(".gitignore", "build/\n*.log\n") + mustWrite("debug.log", "hello log line\n") + return dir +} + +func TestGrepBasic(t *testing.T) { + dir := seedTree(t) + tool := &GrepTool{Root: dir} + res := runSearch(t, tool, map[string]any{"pattern": "hello"}) + txt := resultText(res) + // Matches in tracked files. + if !strings.Contains(txt, "main.go") || !strings.Contains(txt, "util.go") { + t.Errorf("expected go file matches, got %q", txt) + } + // .gitignore'd paths must be skipped. + if strings.Contains(txt, "build/generated.go") { + t.Errorf("ignored dir should be skipped: %q", txt) + } + if strings.Contains(txt, "debug.log") { + t.Errorf("ignored *.log should be skipped: %q", txt) + } +} + +func TestGrepGlobFilter(t *testing.T) { + dir := seedTree(t) + tool := &GrepTool{Root: dir} + res := runSearch(t, tool, map[string]any{"pattern": "hello", "glob": "*.md"}) + txt := resultText(res) + if !strings.Contains(txt, "README.md") { + t.Errorf("expected README match, got %q", txt) + } + if strings.Contains(txt, ".go") { + t.Errorf("glob *.md should exclude .go files: %q", txt) + } +} + +func TestGrepInvalidPattern(t *testing.T) { + dir := seedTree(t) + tool := &GrepTool{Root: dir} + res := runSearch(t, tool, map[string]any{"pattern": "["}) + if !strings.Contains(resultText(res), "invalid pattern") { + t.Errorf("expected invalid-pattern error, got %q", resultText(res)) + } +} + +func TestFindGlob(t *testing.T) { + dir := seedTree(t) + tool := &FindTool{Root: dir} + res := runSearch(t, tool, map[string]any{"glob": "*.go"}) + txt := resultText(res) + if !strings.Contains(txt, "main.go") || !strings.Contains(txt, "sub/deep.go") { + t.Errorf("expected go files, got %q", txt) + } + if strings.Contains(txt, "build/generated.go") { + t.Errorf("ignored dir should be skipped: %q", txt) + } + if strings.Contains(txt, "README.md") { + t.Errorf("*.go should not match README.md: %q", txt) + } +} + +func TestLsDistinguishesFilesAndDirs(t *testing.T) { + dir := seedTree(t) + tool := &LsTool{Root: dir} + res := runSearch(t, tool, map[string]any{}) + txt := resultText(res) + // Directories carry a trailing slash. + if !strings.Contains(txt, "sub/") { + t.Errorf("expected sub/ dir marker, got %q", txt) + } + if !strings.Contains(txt, "main.go") { + t.Errorf("expected main.go file, got %q", txt) + } + details, ok := res.Details.(map[string]any) + if !ok { + t.Fatalf("details missing: %+v", res.Details) + } + if details["files"] == nil || details["dirs"] == nil { + t.Errorf("expected file/dir counts, got %+v", details) + } +} + +func TestLsNotADirectory(t *testing.T) { + dir := seedTree(t) + tool := &LsTool{Root: dir} + res := runSearch(t, tool, map[string]any{"path": "main.go"}) + if !strings.Contains(resultText(res), "not a directory") { + t.Errorf("expected not-a-directory error, got %q", resultText(res)) + } +} + +func TestLsMissing(t *testing.T) { + dir := seedTree(t) + tool := &LsTool{Root: dir} + res := runSearch(t, tool, map[string]any{"path": "nope"}) + if !strings.Contains(resultText(res), "does not exist") { + t.Errorf("expected does-not-exist error, got %q", resultText(res)) + } +} + +func TestSearchPathTraversal(t *testing.T) { + dir := seedTree(t) + for _, tc := range []struct { + name string + tool agentcore.AgentTool + args map[string]any + }{ + {"grep", &GrepTool{Root: dir}, map[string]any{"pattern": "x", "path": "../"}}, + {"find", &FindTool{Root: dir}, map[string]any{"glob": "*", "path": "../"}}, + {"ls", &LsTool{Root: dir}, map[string]any{"path": "../"}}, + } { + t.Run(tc.name, func(t *testing.T) { + res := runSearch(t, tc.tool, tc.args) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Errorf("expected boundary error, got %q", resultText(res)) + } + }) + } +} + +func TestSearchToolModes(t *testing.T) { + for _, tool := range []agentcore.AgentTool{&GrepTool{}, &FindTool{}, &LsTool{}} { + if tool.ExecutionMode() != agentcore.ToolExecutionParallel { + t.Errorf("%s should be parallel", tool.Name()) + } + var schema map[string]any + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Errorf("%s schema not valid JSON: %v", tool.Name(), err) + } + } +} + +func TestGitignoreNegation(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, ".gitignore"), []byte("*.txt\n!keep.txt\n"), 0o644); err != nil { + t.Fatal(err) + } + gi := loadGitignore(dir) + if !gi.ignored("drop.txt", false) { + t.Error("*.txt should be ignored") + } + if gi.ignored("keep.txt", false) { + t.Error("!keep.txt should be re-included") + } +} + +// TestGitignoreMatchModes locks in the three matching modes after the load-time +// precompilation (matchFull / hasSegmentRule): a non-anchored name rule matches +// any path segment (so an ignored dir hides everything beneath it); an anchored +// rule matches only at the root; and a slash-bearing pattern matches the full +// relative path. +func TestGitignoreMatchModes(t *testing.T) { + dir := t.TempDir() + // node_modules: non-anchored → matches any segment (nested too). + // /root.log: anchored → only at repo root. + // a/b.tmp: contains "/" → full-path match. + rules := "node_modules\n/root.log\na/b.tmp\n" + if err := os.WriteFile(filepath.Join(dir, ".gitignore"), []byte(rules), 0o644); err != nil { + t.Fatal(err) + } + gi := loadGitignore(dir) + cases := []struct { + path string + dir bool + want bool + }{ + {"node_modules", true, true}, // segment rule, top level + {"pkg/node_modules", true, true}, // segment rule, nested + {"pkg/node_modules/x/y.js", false, true}, // hidden beneath ignored dir + {"root.log", false, true}, // anchored, at root + {"sub/root.log", false, false}, // anchored must not match nested + {"a/b.tmp", false, true}, // full-path match + {"z/a/b.tmp", false, false}, // full-path rule not anchored elsewhere + {"keep.go", false, false}, // unrelated + } + for _, c := range cases { + if got := gi.ignored(c.path, c.dir); got != c.want { + t.Errorf("ignored(%q, dir=%v) = %v, want %v", c.path, c.dir, got, c.want) + } + } +} diff --git a/pigo/internal/agenttool/todo_tool.go b/pigo/internal/agenttool/todo_tool.go new file mode 100644 index 0000000..db528ce --- /dev/null +++ b/pigo/internal/agenttool/todo_tool.go @@ -0,0 +1,186 @@ +// This file implements the todo tool (US-011, #127): a structured task list the +// model uses to plan and track multi-step work, with progress visible to the +// user. Unlike the file tools this one is stateful — the written list lives in a +// per-session TodoStore the tool holds, so a later write replaces the plan and +// the REPL can render the current progress after each update. +// +// pi itself has no such tool; this mirrors Claude Code's TodoWrite: the model +// submits the WHOLE list each call (not incremental edits), each item carries a +// content string and a status of pending | in_progress | completed. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// TodoStatus is the lifecycle state of a single todo item. +type TodoStatus string + +const ( + // TodoPending is a task not yet started. + TodoPending TodoStatus = "pending" + // TodoInProgress is the task currently being worked on. + TodoInProgress TodoStatus = "in_progress" + // TodoCompleted is a finished task. + TodoCompleted TodoStatus = "completed" +) + +// validTodoStatus reports whether s is one of the three accepted statuses. +func validTodoStatus(s TodoStatus) bool { + switch s { + case TodoPending, TodoInProgress, TodoCompleted: + return true + default: + return false + } +} + +// TodoItem is one entry in the task list. +type TodoItem struct { + // Content is the human-readable task description. + Content string `json:"content"` + // Status is the item's lifecycle state. + Status TodoStatus `json:"status"` +} + +// TodoStore holds the current task list for a session. It is safe for concurrent +// use so the tool (which may run in a batch) and the REPL renderer can touch it +// without racing. A single store is shared for a session's lifetime. +type TodoStore struct { + mu sync.RWMutex + items []TodoItem +} + +// NewTodoStore returns an empty store. +func NewTodoStore() *TodoStore { return &TodoStore{} } + +// Set replaces the whole list with items (a copy, so the caller's slice can be +// reused). +func (s *TodoStore) Set(items []TodoItem) { + s.mu.Lock() + defer s.mu.Unlock() + s.items = append(s.items[:0:0], items...) +} + +// Snapshot returns a copy of the current list, safe to read without holding the +// lock. +func (s *TodoStore) Snapshot() []TodoItem { + s.mu.RLock() + defer s.mu.RUnlock() + return append([]TodoItem(nil), s.items...) +} + +// TodoTool is the stateful todo-list tool. It writes the submitted list into +// Store, replacing any previous list, and returns a rendered progress view. +type TodoTool struct { + // Store holds the session task list. Must be non-nil; NewTodoStore builds one. + Store *TodoStore +} + +// todoToolArgs is the decoded argument shape: the full task list to store. +type todoToolArgs struct { + Todos []TodoItem `json:"todos"` +} + +// Name implements AgentTool. +func (t *TodoTool) Name() string { return "todo" } + +// Description implements AgentTool. +func (t *TodoTool) Description() string { + return "Record and update a structured task list to plan and track multi-step " + + "work. Submit the ENTIRE list every call; it replaces the previous list. " + + "Each item has a content string and a status of pending, in_progress, or " + + "completed. Keep exactly one item in_progress at a time and mark items " + + "completed as soon as they are done." +} + +// Schema implements AgentTool. +func (t *TodoTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "todos": { + "type": "array", + "description": "The full task list, replacing any previous list.", + "items": { + "type": "object", + "properties": { + "content": {"type": "string", "description": "Task description."}, + "status": {"type": "string", "enum": ["pending", "in_progress", "completed"], "description": "Task lifecycle state."} + }, + "required": ["content", "status"], + "additionalProperties": false + } + } + }, + "required": ["todos"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Updating the shared list mutates session +// state → sequential so a batch cannot interleave two list writes. +func (t *TodoTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// Execute implements AgentTool. It validates every item's status, stores the +// list, and returns the rendered progress as the result content. Invalid input +// degrades to an error result (matching the file tools) rather than a Go error. +func (t *TodoTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[todoToolArgs](args, "todo") + if bad != nil { + return *bad, nil + } + for i, it := range a.Todos { + if strings.TrimSpace(it.Content) == "" { + return errorResult(fmt.Sprintf("todo: item %d has empty content", i+1)), nil + } + if !validTodoStatus(it.Status) { + return errorResult(fmt.Sprintf("todo: item %d has invalid status %q (want pending|in_progress|completed)", i+1, it.Status)), nil + } + } + + if t.Store == nil { + t.Store = NewTodoStore() + } + t.Store.Set(a.Todos) + + rendered := RenderTodoList(a.Todos) + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(rendered)}, + Details: map[string]any{"todos": a.Todos}, + }, nil +} + +// RenderTodoList renders items as a checkbox progress block, one line per task, +// with a trailing summary count. An empty list renders as a single "(no tasks)" +// line so an intentional clear is still visible. The marks are: [ ] pending, +// [~] in_progress, [x] completed. +func RenderTodoList(items []TodoItem) string { + if len(items) == 0 { + return "Todos: (no tasks)" + } + var b strings.Builder + done := 0 + b.WriteString("Todos:") + for _, it := range items { + mark := " " + switch it.Status { + case TodoInProgress: + mark = "~" + case TodoCompleted: + mark = "x" + done++ + } + fmt.Fprintf(&b, "\n [%s] %s", mark, it.Content) + } + fmt.Fprintf(&b, "\n(%d/%d completed)", done, len(items)) + return b.String() +} diff --git a/pigo/internal/agenttool/todo_tool_test.go b/pigo/internal/agenttool/todo_tool_test.go new file mode 100644 index 0000000..545202e --- /dev/null +++ b/pigo/internal/agenttool/todo_tool_test.go @@ -0,0 +1,143 @@ +// Tests for the todo tool (US-011, #127): registration/validation, status +// transitions across successive writes, and the rendered progress view. +package agenttool + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// execTodo runs the tool with the given JSON args and returns the result. +func execTodo(t *testing.T, tool *TodoTool, args string) agentcore.AgentToolResult { + t.Helper() + res, err := tool.Execute(context.Background(), "call-1", json.RawMessage(args), nil) + if err != nil { + t.Fatalf("Execute returned Go error: %v", err) + } + return res +} + +// TestTodoToolRegisters checks the tool registers cleanly (valid schema) and is +// retrievable — the "registered in the agenttool registry" acceptance criterion. +func TestTodoToolRegisters(t *testing.T) { + reg := NewToolRegistry() + tool := &TodoTool{Store: NewTodoStore()} + if err := reg.Register(tool); err != nil { + t.Fatalf("Register: %v", err) + } + got, ok := reg.Get("todo") + if !ok { + t.Fatal("todo tool not found after Register") + } + if got.Name() != "todo" { + t.Errorf("Name = %q, want todo", got.Name()) + } +} + +// TestTodoToolStoresList checks a write lands in the session store. +func TestTodoToolStoresList(t *testing.T) { + store := NewTodoStore() + tool := &TodoTool{Store: store} + execTodo(t, tool, `{"todos":[ + {"content":"first","status":"in_progress"}, + {"content":"second","status":"pending"} + ]}`) + + items := store.Snapshot() + if len(items) != 2 { + t.Fatalf("Snapshot len = %d, want 2", len(items)) + } + if items[0].Content != "first" || items[0].Status != TodoInProgress { + t.Errorf("item 0 = %+v", items[0]) + } + if items[1].Status != TodoPending { + t.Errorf("item 1 status = %q, want pending", items[1].Status) + } +} + +// TestTodoToolStatusTransition checks a later write replaces the previous list, +// reflecting a status flow pending → in_progress → completed. +func TestTodoToolStatusTransition(t *testing.T) { + store := NewTodoStore() + tool := &TodoTool{Store: store} + + execTodo(t, tool, `{"todos":[{"content":"build feature","status":"pending"}]}`) + if s := store.Snapshot()[0].Status; s != TodoPending { + t.Fatalf("after write 1 status = %q, want pending", s) + } + + execTodo(t, tool, `{"todos":[{"content":"build feature","status":"in_progress"}]}`) + if s := store.Snapshot()[0].Status; s != TodoInProgress { + t.Fatalf("after write 2 status = %q, want in_progress", s) + } + + execTodo(t, tool, `{"todos":[{"content":"build feature","status":"completed"}]}`) + items := store.Snapshot() + if len(items) != 1 { + t.Fatalf("after write 3 len = %d, want 1", len(items)) + } + if items[0].Status != TodoCompleted { + t.Errorf("after write 3 status = %q, want completed", items[0].Status) + } +} + +// TestTodoToolRejectsInvalidStatus checks an unknown status degrades to an error +// result and does not mutate the store. +func TestTodoToolRejectsInvalidStatus(t *testing.T) { + store := NewTodoStore() + tool := &TodoTool{Store: store} + res := execTodo(t, tool, `{"todos":[{"content":"x","status":"done"}]}`) + if !strings.Contains(agentcore.ContentToText(res.Content), "invalid status") { + t.Errorf("expected invalid-status error, got %q", agentcore.ContentToText(res.Content)) + } + if len(store.Snapshot()) != 0 { + t.Error("store mutated despite invalid input") + } +} + +// TestTodoToolRejectsEmptyContent checks a blank content is rejected. +func TestTodoToolRejectsEmptyContent(t *testing.T) { + tool := &TodoTool{Store: NewTodoStore()} + res := execTodo(t, tool, `{"todos":[{"content":" ","status":"pending"}]}`) + if !strings.Contains(agentcore.ContentToText(res.Content), "empty content") { + t.Errorf("expected empty-content error, got %q", agentcore.ContentToText(res.Content)) + } +} + +// TestRenderTodoList checks the rendered progress view: marks per status and a +// completion count. +func TestRenderTodoList(t *testing.T) { + out := RenderTodoList([]TodoItem{ + {Content: "alpha", Status: TodoCompleted}, + {Content: "beta", Status: TodoInProgress}, + {Content: "gamma", Status: TodoPending}, + }) + for _, want := range []string{"[x] alpha", "[~] beta", "[ ] gamma", "(1/3 completed)"} { + if !strings.Contains(out, want) { + t.Errorf("render missing %q in:\n%s", want, out) + } + } +} + +// TestRenderTodoListEmpty checks an empty list renders visibly (an intentional +// clear should still show). +func TestRenderTodoListEmpty(t *testing.T) { + if out := RenderTodoList(nil); !strings.Contains(out, "no tasks") { + t.Errorf("empty render = %q, want a (no tasks) marker", out) + } +} + +// TestTodoToolResultRenders checks Execute returns the rendered list as content +// so the REPL has something to display. +func TestTodoToolResultRenders(t *testing.T) { + tool := &TodoTool{Store: NewTodoStore()} + res := execTodo(t, tool, `{"todos":[{"content":"do it","status":"completed"}]}`) + text := agentcore.ContentToText(res.Content) + if !strings.Contains(text, "[x] do it") || !strings.Contains(text, "(1/1 completed)") { + t.Errorf("result content = %q", text) + } +} diff --git a/pigo/internal/agenttool/tool_executor.go b/pigo/internal/agenttool/tool_executor.go new file mode 100644 index 0000000..09797b3 --- /dev/null +++ b/pigo/internal/agenttool/tool_executor.go @@ -0,0 +1,341 @@ +// This file implements the three-phase tool execution (US-004): prepare → +// execute → finalize, with the beforeToolCall / afterToolCall hooks. It mirrors +// pi's agent-loop tool handling: a tool call is looked up in the registry, its +// arguments are (optionally) prepared and schema-validated, the beforeToolCall +// hook may block it, the tool runs (streaming partial updates), and the +// afterToolCall hook may override the result field-by-field (no deep merge). +// +// Every failure mode (unknown tool, validation failure, block, abort, tool +// error/panic) is turned into an error tool result rather than a Go error, so +// the loop always has a ToolResultMessage to feed back to the model. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// toolResultMaxBytes is the executor-layer budget for a single tool result's +// combined text, applied uniformly to EVERY tool right before its result enters +// the AgentToolResult / message list. Individual tools also impose their own, +// stricter inner caps (read: readToolMaxLines, search: searchMaxResults, +// webfetch: webFetchMaxBytes, bash: bashMaxOutputBytes); those still run first +// and clip a tool below this outer budget. This budget is the last line of +// defense so a tool with no (or a looser) inner cap cannot blow the model's +// context. Override per-executor via ToolExecutorConfig.MaxResultBytes. +const toolResultMaxBytes = 100_000 + +// ToolExecutorConfig holds the registry and the optional per-phase hooks. Every +// hook is optional (nil = default behavior). +type ToolExecutorConfig struct { + Registry *ToolRegistry + PrepareArguments agentcore.PrepareArgumentsFunc + BeforeToolCall agentcore.BeforeToolCallFunc + AfterToolCall agentcore.AfterToolCallFunc + // MaxResultBytes overrides the executor-layer per-result text budget. Zero + // (the default) uses toolResultMaxBytes; a negative value disables the + // budget entirely. + MaxResultBytes int + // MaxToolRetries overrides the number of RETRIES for a transient tool error + // (see isRetryableToolError). Zero (the default) uses maxToolRetries; a + // negative value disables retrying (a single attempt). Mirrors the + // MaxResultBytes sentinel convention. + MaxToolRetries int +} + +// executeToolCall runs one tool call through prepare → execute → finalize and +// returns the resulting ToolResultMessage plus whether the batch should +// terminate. emit may be nil (no events). It never returns a Go error: every +// failure is encoded into the returned message with IsError=true. +func executeToolCall(ctx context.Context, cfg ToolExecutorConfig, call agentcore.AgentToolCall, emit agentcore.EmitFunc) (agentcore.ToolResultMessage, bool) { + // 1. prepare: lookup, prepareArguments, validate, beforeToolCall. + tool, args, prep, isError := prepareToolCall(ctx, cfg, call) + if prep != nil { + // Prepare short-circuited (unknown tool / prepare error / validation / + // block / abort): finalize the error result without executing. + return finalizeToolCall(ctx, cfg, call, *prep, isError, emit) + } + + // 2. execute. + if emit != nil { + if err := emit(ctx, agentcore.ToolExecutionStartEvent{ToolCallID: call.ID, ToolName: call.Name, Args: args}); err != nil { + return errorToolResult(call, "aborted before execution: "+err.Error()), false + } + } + result, isError := runToolWithRetry(ctx, cfg, tool, call, args, emit) + + // 3. finalize: afterToolCall overrides. + return finalizeToolCall(ctx, cfg, call, result, isError, emit) +} + +// prepareToolCall performs the prepare phase. On success it returns the tool and +// the (possibly rewritten) arguments with a nil result. On any short-circuit it +// returns a non-nil *AgentToolResult and the isError flag. +func prepareToolCall(ctx context.Context, cfg ToolExecutorConfig, call agentcore.AgentToolCall) (agentcore.AgentTool, json.RawMessage, *agentcore.AgentToolResult, bool) { + if ctx.Err() != nil { + r := errorResult(fmt.Sprintf("tool %q aborted before execution", call.Name)) + return nil, nil, &r, true + } + + // Registry lookup. + tool, ok := cfg.Registry.Get(call.Name) + if !ok { + r := errorResult(fmt.Sprintf("unknown tool %q", call.Name)) + return nil, nil, &r, true + } + + // prepareArguments (optional). + args := call.Arguments + if cfg.PrepareArguments != nil { + prepared, err := cfg.PrepareArguments(ctx, call.Name, args) + if err != nil { + r := errorResult(fmt.Sprintf("prepareArguments for %q failed: %v", call.Name, err)) + return nil, nil, &r, true + } + args = prepared + } + + // JSON Schema validation. + if errs := cfg.Registry.Validate(call.Name, args); len(errs) > 0 { + r := ValidationErrorResult(call.Name, errs) + return nil, nil, &r, true + } + + // beforeToolCall hook (may block or rewrite arguments). + if cfg.BeforeToolCall != nil { + if dec := cfg.BeforeToolCall(ctx, agentcore.AgentToolCall{ID: call.ID, Name: call.Name, Arguments: args}); dec != nil { + if dec.Block { + r := agentcore.AgentToolResult{} + if dec.Content != nil { + r.Content = *dec.Content + } else { + r.Content = agentcore.ContentList{agentcore.NewTextContent(fmt.Sprintf("tool %q blocked by beforeToolCall", call.Name))} + } + if dec.Details != nil { + r.Details = *dec.Details + } + return nil, nil, &r, true + } + // Argument rewrite (PreToolUse updatedInput): replace and re-validate + // so a hook cannot smuggle schema-invalid args past the tool. + if len(dec.UpdatedInput) > 0 { + args = dec.UpdatedInput + if errs := cfg.Registry.Validate(call.Name, args); len(errs) > 0 { + r := ValidationErrorResult(call.Name, errs) + return nil, nil, &r, true + } + } + } + } + + return tool, args, nil, false +} + +// runToolWithRetry wraps runTool with the classified, bounded retry policy at +// the single tool-execution seam so EVERY tool gets uniform resilience. It only +// retries when runTool surfaces a non-nil Go error (transport/agent-error path) +// AND isRetryableToolError says that error is transient; a (result, nil) is +// done regardless of the result's IsError flag (a tool's own terminal result is +// never retried). Retries are capped by toolRetryCap and separated by a small +// backoff. Context cancellation short-circuits immediately: a cancelled/expired +// outer ctx is never retried. +func runToolWithRetry(ctx context.Context, cfg ToolExecutorConfig, tool agentcore.AgentTool, call agentcore.AgentToolCall, args json.RawMessage, emit agentcore.EmitFunc) (agentcore.AgentToolResult, bool) { + retryCap := toolRetryCap(cfg.MaxToolRetries) + + var lastResult agentcore.AgentToolResult + var lastIsError bool + for attempt := 0; attempt <= retryCap; attempt++ { + result, err, isError := runTool(ctx, tool, call, args, emit) + if err == nil { + // Execute returned (result, nil): terminal success regardless of + // the result's own IsError flag. Done, no retry. + return result, isError + } + + lastResult, lastIsError = result, isError + + // Do not retry if the outer context is done (Canceled or its deadline + // has passed) — a dead context means stop. + if ctx.Err() != nil { + break + } + // Only transient errors are retried, and only if we have budget left. + if attempt >= retryCap || !isRetryableToolError(err) { + break + } + // Small backoff; abort the wait early if ctx dies mid-sleep. + if !waitToolRetryBackoff(ctx, attempt) { + break + } + } + return lastResult, lastIsError +} + +// runTool executes the tool, recovering a panic into an error. It returns the +// shaped error result, the raw error (nil on success), and the isError flag. +// The raw error is surfaced so the caller can classify it for retry; on success +// err is nil even if the result itself carries IsError semantics. +func runTool(ctx context.Context, tool agentcore.AgentTool, call agentcore.AgentToolCall, args json.RawMessage, emit agentcore.EmitFunc) (result agentcore.AgentToolResult, rawErr error, isError bool) { + defer func() { + if r := recover(); r != nil { + result = errorResult(fmt.Sprintf("tool %q panicked: %v", call.Name, r)) + rawErr = toolPanic{value: r} + isError = true + } + }() + + onUpdate := func(partial agentcore.AgentToolResult) { + if emit == nil { + return + } + _ = emit(ctx, agentcore.ToolExecutionUpdateEvent{ToolCallID: call.ID, ToolName: call.Name, PartialResult: partial}) + } + + res, err := tool.Execute(ctx, call.ID, args, onUpdate) + if err != nil { + return errorResult(fmt.Sprintf("tool %q failed: %v", call.Name, err)), err, true + } + return res, nil, false +} + +// finalizeToolCall applies the afterToolCall hook (field-level override, no deep +// merge), emits the tool_execution_end event, and builds the ToolResultMessage. +// It returns the message and whether this result requests termination. +func finalizeToolCall(ctx context.Context, cfg ToolExecutorConfig, call agentcore.AgentToolCall, result agentcore.AgentToolResult, isError bool, emit agentcore.EmitFunc) (agentcore.ToolResultMessage, bool) { + if cfg.AfterToolCall != nil { + if ov := cfg.AfterToolCall(ctx, call, result, isError); ov != nil { + if ov.Content != nil { + result.Content = *ov.Content + } + if ov.Details != nil { + result.Details = *ov.Details + } + if ov.Terminate != nil { + result.Terminate = ov.Terminate + } + if ov.IsError != nil { + isError = *ov.IsError + } + } + } + + // Result-shaping seam: every tool's output funnels through here before it + // becomes a ToolResultMessage, so this is the single point where the + // executor-layer byte budget is enforced uniformly for ALL tools. + result.Content = clipToolResultContent(result.Content, cfg.MaxResultBytes) + + if emit != nil { + _ = emit(ctx, agentcore.ToolExecutionEndEvent{ToolCallID: call.ID, ToolName: call.Name, Result: result, IsError: isError}) + } + + terminate := result.Terminate != nil && *result.Terminate + return agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, + ToolCallID: call.ID, + ToolName: call.Name, + Content: result.Content, + Details: result.Details, + IsError: isError, + }, terminate +} + +// errorResult builds an error AgentToolResult carrying a single text block. +func errorResult(msg string) agentcore.AgentToolResult { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(msg)}} +} + +// clipToolResultContent enforces the executor-layer byte budget on a tool +// result's text, uniformly for every tool. budget<=0 with the sentinel meaning: +// 0 => toolResultMaxBytes default, <0 => disabled. Non-text blocks (e.g. images) +// pass through untouched and keep their order; the combined text of all text +// blocks is measured against the budget and, when over, collapsed into a single +// truncated text block via truncateToBudget (head + "[truncated N bytes]" + +// tail, matching the bash idiom). Per-tool inner caps have already run, so this +// only bites when a tool's own cap is looser or absent. +func clipToolResultContent(content agentcore.ContentList, cfgMax int) agentcore.ContentList { + budget := cfgMax + if budget == 0 { + budget = toolResultMaxBytes + } + if budget < 0 { + return content + } + + total := 0 + textBlocks := 0 + for _, c := range content { + if t, ok := c.(agentcore.TextContent); ok { + total += len(t.Text) + textBlocks++ + } + } + if textBlocks == 0 || total <= budget { + return content + } + + // Over budget: gather all text (in order) and non-text blocks separately, + // then emit the non-text blocks followed by one truncated text block. + var sb strings.Builder + out := make(agentcore.ContentList, 0, len(content)) + for _, c := range content { + if t, ok := c.(agentcore.TextContent); ok { + sb.WriteString(t.Text) + continue + } + out = append(out, c) + } + out = append(out, agentcore.NewTextContent(truncateToBudget(sb.String(), budget))) + return out +} + +// truncateToBudget caps s at budget bytes. When s is longer it keeps a head and +// a tail preview (split evenly) joined by a "[truncated N bytes]" marker, so +// both the start and the end of the text survive. Cut points are pulled back to +// UTF-8 rune boundaries so no partial rune is emitted; N counts the raw bytes +// dropped from the middle. This is the single shared truncation idiom reused by +// both the bash tool's inner cap and the executor-layer budget. +func truncateToBudget(s string, budget int) string { + if budget <= 0 || len(s) <= budget { + return s + } + half := budget / 2 + head := trimUTF8Prefix(s[:half]) + tail := trimUTF8Suffix(s[len(s)-half:]) + removed := len(s) - len(head) - len(tail) + return head + fmt.Sprintf("\n[truncated %d bytes]\n", removed) + tail +} + +// decodeArgs unmarshals a tool's JSON arguments into T. On failure it returns an +// error result already shaped as ": invalid arguments: ...", so a tool's +// Execute can decode and bail in one line: +// +// a, bad := decodeArgs[readToolArgs](args, "read") +// if bad != nil { +// return *bad, nil +// } +// +// The ok flag distinguishes the failure case without comparing the zero value. +func decodeArgs[T any](args json.RawMessage, tool string) (T, *agentcore.AgentToolResult) { + var a T + if err := json.Unmarshal(args, &a); err != nil { + res := errorResult(fmt.Sprintf("%s: invalid arguments: %v", tool, err)) + return a, &res + } + return a, nil +} + +// errorToolResult builds an error ToolResultMessage directly (used when a call +// is aborted outside the normal finalize path). +func errorToolResult(call agentcore.AgentToolCall, msg string) agentcore.ToolResultMessage { + return agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, + ToolCallID: call.ID, + ToolName: call.Name, + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + IsError: true, + } +} diff --git a/pigo/internal/agenttool/tool_executor_test.go b/pigo/internal/agenttool/tool_executor_test.go new file mode 100644 index 0000000..9838b6b --- /dev/null +++ b/pigo/internal/agenttool/tool_executor_test.go @@ -0,0 +1,465 @@ +package agenttool + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "os" + "strings" + "sync/atomic" + "syscall" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// execTool is a configurable AgentTool for executor tests. +type execTool struct { + name string + schema string + run func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) + mode agentcore.ToolExecutionMode +} + +func (t execTool) Name() string { return t.name } +func (t execTool) Description() string { return "exec" } +func (t execTool) Schema() json.RawMessage { + if t.schema == "" { + return nil + } + return json.RawMessage(t.schema) +} +func (t execTool) ExecutionMode() agentcore.ToolExecutionMode { + if t.mode == "" { + return agentcore.ToolExecutionParallel + } + return t.mode +} +func (t execTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return t.run(ctx, id, args, onUpdate) +} + +func newExecCfg(t *testing.T, tool agentcore.AgentTool) ToolExecutorConfig { + t.Helper() + r := NewToolRegistry() + if err := r.Register(tool); err != nil { + t.Fatalf("register: %v", err) + } + return ToolExecutorConfig{Registry: r} +} + +func textOf(msg agentcore.ToolResultMessage) string { + if len(msg.Content) == 0 { + return "" + } + if tc, ok := msg.Content[0].(agentcore.TextContent); ok { + return tc.Text + } + return "" +} + +func TestExecutorNormal(t *testing.T) { + tool := execTool{name: "echo", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("done")}}, nil + }} + cfg := newExecCfg(t, tool) + + var events []agentcore.AgentEvent + emit := func(ctx context.Context, ev agentcore.AgentEvent) error { events = append(events, ev); return nil } + msg, term := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "echo"}, emit) + + if msg.IsError || textOf(msg) != "done" { + t.Fatalf("normal result wrong: %+v", msg) + } + if term { + t.Error("normal result should not terminate") + } + wantKinds := []string{agentcore.EventToolExecutionStart, agentcore.EventToolExecutionEnd} + if len(events) != 2 || events[0].EventType() != wantKinds[0] || events[1].EventType() != wantKinds[1] { + t.Errorf("events wrong: %+v", events) + } +} + +func TestExecutorUnknownTool(t *testing.T) { + cfg := ToolExecutorConfig{Registry: NewToolRegistry()} + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "ghost"}, nil) + if !msg.IsError { + t.Fatalf("unknown tool should be error result: %+v", msg) + } +} + +func TestExecutorValidationFailure(t *testing.T) { + schema := `{"type":"object","properties":{"n":{"type":"integer"}},"required":["n"],"additionalProperties":false}` + tool := execTool{name: "need", schema: schema, run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + t.Fatal("execute must not run on validation failure") + return agentcore.AgentToolResult{}, nil + }} + cfg := newExecCfg(t, tool) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "need", Arguments: json.RawMessage(`{}`)}, nil) + if !msg.IsError { + t.Fatalf("validation failure should be error result: %+v", msg) + } +} + +func TestExecutorBlock(t *testing.T) { + tool := execTool{name: "echo", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + t.Fatal("execute must not run when blocked") + return agentcore.AgentToolResult{}, nil + }} + cfg := newExecCfg(t, tool) + cfg.BeforeToolCall = func(ctx context.Context, call agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { + return &agentcore.BeforeToolCallDecision{Block: true} + } + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "echo"}, nil) + if !msg.IsError { + t.Fatalf("blocked call should be error result: %+v", msg) + } +} + +func TestExecutorToolError(t *testing.T) { + tool := execTool{name: "boom", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{}, errors.New("kaboom") + }} + cfg := newExecCfg(t, tool) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "boom"}, nil) + if !msg.IsError { + t.Fatalf("tool error should be error result: %+v", msg) + } +} + +func TestExecutorPanicRecovered(t *testing.T) { + tool := execTool{name: "panic", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + panic("oops") + }} + cfg := newExecCfg(t, tool) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "panic"}, nil) + if !msg.IsError { + t.Fatalf("panic should be recovered into error result: %+v", msg) + } +} + +func TestExecutorAfterToolCallOverride(t *testing.T) { + tool := execTool{name: "echo", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("orig")}}, nil + }} + cfg := newExecCfg(t, tool) + newContent := agentcore.ContentList{agentcore.NewTextContent("overridden")} + isErr := true + term := true + cfg.AfterToolCall = func(ctx context.Context, call agentcore.AgentToolCall, result agentcore.AgentToolResult, isError bool) *agentcore.AfterToolCallResult { + return &agentcore.AfterToolCallResult{Content: &newContent, IsError: &isErr, Terminate: &term} + } + msg, terminate := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "echo"}, nil) + if textOf(msg) != "overridden" { + t.Errorf("content override failed: %q", textOf(msg)) + } + if !msg.IsError { + t.Error("isError override failed") + } + if !terminate { + t.Error("terminate override failed") + } +} + +func TestExecutorUpdateCallback(t *testing.T) { + tool := execTool{name: "stream", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + onUpdate(agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("partial")}}) + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("final")}}, nil + }} + cfg := newExecCfg(t, tool) + var updates int + emit := func(ctx context.Context, ev agentcore.AgentEvent) error { + if ev.EventType() == agentcore.EventToolExecutionUpdate { + updates++ + } + return nil + } + executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "stream"}, emit) + if updates != 1 { + t.Errorf("expected 1 update event, got %d", updates) + } +} + +func TestExecutorAbortedContext(t *testing.T) { + tool := execTool{name: "echo", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + t.Fatal("execute must not run when context already cancelled") + return agentcore.AgentToolResult{}, nil + }} + cfg := newExecCfg(t, tool) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + msg, _ := executeToolCall(ctx, cfg, agentcore.AgentToolCall{ID: "1", Name: "echo"}, nil) + if !msg.IsError { + t.Fatalf("aborted call should be error result: %+v", msg) + } +} + +// TestExecutorResultBudget proves the executor-layer byte budget applies to any +// tool: a stub tool emitting output larger than the budget gets its result text +// truncated with an accurate "[truncated N bytes]" marker, while a small output +// is left untouched. +func TestExecutorResultBudget(t *testing.T) { + const budget = 1000 + big := strings.Repeat("A", budget) + strings.Repeat("B", budget) // 2*budget bytes + tool := execTool{name: "flood", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(big)}}, nil + }} + cfg := newExecCfg(t, tool) + cfg.MaxResultBytes = budget + + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "flood"}, nil) + got := textOf(msg) + if len(got) >= len(big) { + t.Fatalf("output not truncated: len=%d, original=%d", len(got), len(big)) + } + half := budget / 2 + removed := len(big) - 2*half + marker := fmt.Sprintf("\n[truncated %d bytes]\n", removed) + if !strings.Contains(got, marker) { + t.Fatalf("missing/incorrect truncation marker %q in output %q", marker, got) + } + want := big[:half] + marker + big[len(big)-half:] + if got != want { + t.Fatalf("truncated output mismatch:\n got=%q\nwant=%q", got, want) + } +} + +// TestExecutorResultBudgetSmallOutputUntouched proves outputs within budget are +// passed through verbatim. +func TestExecutorResultBudgetSmallOutputUntouched(t *testing.T) { + small := "just a little output" + tool := execTool{name: "tiny", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(small)}}, nil + }} + cfg := newExecCfg(t, tool) + cfg.MaxResultBytes = 1000 + + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "tiny"}, nil) + if got := textOf(msg); got != small { + t.Fatalf("small output altered: got=%q want=%q", got, small) + } +} + +// TestExecutorResultBudgetDefault proves the default (zero MaxResultBytes) uses +// toolResultMaxBytes and truncates output beyond it. +func TestExecutorResultBudgetDefault(t *testing.T) { + big := strings.Repeat("x", toolResultMaxBytes+5000) + tool := execTool{name: "flood", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(big)}}, nil + }} + cfg := newExecCfg(t, tool) // MaxResultBytes == 0 -> default + + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "flood"}, nil) + got := textOf(msg) + if !strings.Contains(got, "[truncated ") { + t.Fatalf("default budget did not truncate: len=%d", len(got)) + } + if len(got) > toolResultMaxBytes+64 { + t.Fatalf("default-truncated output too large: %d", len(got)) + } +} + +// TestExecutorResultBudgetDisabled proves a negative MaxResultBytes disables the +// budget entirely. +func TestExecutorResultBudgetDisabled(t *testing.T) { + big := strings.Repeat("y", toolResultMaxBytes*2) + tool := execTool{name: "flood", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(big)}}, nil + }} + cfg := newExecCfg(t, tool) + cfg.MaxResultBytes = -1 + + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "flood"}, nil) + if got := textOf(msg); got != big { + t.Fatalf("disabled budget altered output: len=%d want=%d", len(got), len(big)) + } +} + +// --- Tool-execution retry (node #252) --------------------------------------- + +// countingTool returns a transient error for its first failN attempts, then +// succeeds; if failN < 0 it always fails. It records how many times Execute ran. +type retryStub struct { + failN int // number of leading failures before success; <0 = always fail + err error // error to return on a failing attempt + attempts int32 +} + +func (s *retryStub) run(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + n := atomic.AddInt32(&s.attempts, 1) + if s.failN < 0 || int(n) <= s.failN { + return agentcore.AgentToolResult{}, s.err + } + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, nil +} + +func newRetryCfg(t *testing.T, name string, s *retryStub) ToolExecutorConfig { + t.Helper() + return newExecCfg(t, execTool{name: name, run: s.run}) +} + +func TestExecutorRetryTransientThenSuccess(t *testing.T) { + // Fails 2 times with a transient error, then succeeds. With the default cap + // (2 retries = 3 attempts) this should ultimately succeed on attempt 3. + s := &retryStub{failN: 2, err: syscall.ECONNRESET} + cfg := newRetryCfg(t, "flaky", s) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "flaky"}, nil) + if msg.IsError { + t.Fatalf("expected eventual success, got error: %q", textOf(msg)) + } + if got := atomic.LoadInt32(&s.attempts); got != 3 { + t.Fatalf("expected 3 attempts (2 retries), got %d", got) + } +} + +func TestExecutorRetryCapExhausted(t *testing.T) { + // Always fails with a transient error: must stop after maxToolRetries+1 + // attempts (default cap) and give up with an error result. + s := &retryStub{failN: -1, err: syscall.ETIMEDOUT} + cfg := newRetryCfg(t, "always", s) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "always"}, nil) + if !msg.IsError { + t.Fatalf("expected error result after exhausting retries: %+v", msg) + } + want := int32(maxToolRetries + 1) + if got := atomic.LoadInt32(&s.attempts); got != want { + t.Fatalf("expected %d attempts, got %d", want, got) + } +} + +func TestExecutorRetryTerminalNoRetry(t *testing.T) { + // A terminal (non-transient) error must be tried exactly once. + s := &retryStub{failN: -1, err: errors.New("invalid argument: bad")} + cfg := newRetryCfg(t, "terminal", s) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "terminal"}, nil) + if !msg.IsError { + t.Fatalf("expected error result: %+v", msg) + } + if got := atomic.LoadInt32(&s.attempts); got != 1 { + t.Fatalf("terminal error must not retry: got %d attempts", got) + } +} + +func TestExecutorRetryDisabled(t *testing.T) { + // MaxToolRetries < 0 disables retry even for a transient error. + s := &retryStub{failN: -1, err: syscall.ECONNRESET} + cfg := newRetryCfg(t, "notretry", s) + cfg.MaxToolRetries = -1 + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "notretry"}, nil) + if !msg.IsError { + t.Fatalf("expected error result: %+v", msg) + } + if got := atomic.LoadInt32(&s.attempts); got != 1 { + t.Fatalf("disabled retry must try once: got %d attempts", got) + } +} + +func TestExecutorRetryCustomCap(t *testing.T) { + // A custom positive cap is honored: always-failing transient error stops at + // cap+1 attempts. + s := &retryStub{failN: -1, err: syscall.EAGAIN} + cfg := newRetryCfg(t, "custom", s) + cfg.MaxToolRetries = 4 + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "custom"}, nil) + if !msg.IsError { + t.Fatalf("expected error result: %+v", msg) + } + if got := atomic.LoadInt32(&s.attempts); got != 5 { + t.Fatalf("expected 5 attempts (cap 4), got %d", got) + } +} + +func TestExecutorRetryCanceledContextNoRetry(t *testing.T) { + // context.Canceled surfaced by the tool must never be retried. + s := &retryStub{failN: -1, err: context.Canceled} + cfg := newRetryCfg(t, "cancel", s) + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "cancel"}, nil) + if !msg.IsError { + t.Fatalf("expected error result: %+v", msg) + } + if got := atomic.LoadInt32(&s.attempts); got != 1 { + t.Fatalf("context.Canceled must not retry: got %d attempts", got) + } +} + +func TestExecutorRetryStopsWhenOuterCtxCanceled(t *testing.T) { + // If the outer ctx is cancelled during execution, retries stop even though + // the returned error is transient. + ctx, cancel := context.WithCancel(context.Background()) + s := &retryStub{failN: -1, err: syscall.ECONNRESET} + tool := execTool{name: "abortmid", run: func(c context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + atomic.AddInt32(&s.attempts, 1) + cancel() // outer ctx dies after the first attempt + return agentcore.AgentToolResult{}, syscall.ECONNRESET + }} + cfg := newExecCfg(t, tool) + msg, _ := executeToolCall(ctx, cfg, agentcore.AgentToolCall{ID: "1", Name: "abortmid"}, nil) + if !msg.IsError { + t.Fatalf("expected error result: %+v", msg) + } + if got := atomic.LoadInt32(&s.attempts); got != 1 { + t.Fatalf("cancelled outer ctx must stop retry: got %d attempts", got) + } +} + +func TestIsRetryableToolError(t *testing.T) { + transient := []error{ + syscall.ETIMEDOUT, + syscall.ECONNRESET, + syscall.EAGAIN, + context.DeadlineExceeded, + os.ErrDeadlineExceeded, + fmt.Errorf("dial tcp: %w", syscall.ECONNRESET), + errors.New("connection refused"), + errors.New("resource temporarily unavailable"), + errors.New("read: i/o timeout"), + &net.DNSError{IsTimeout: true}, + } + for _, err := range transient { + if !isRetryableToolError(err) { + t.Errorf("expected transient (retryable): %v", err) + } + } + terminal := []error{ + nil, + context.Canceled, + fmt.Errorf("wrapped: %w", context.Canceled), + errors.New("file not found"), + errors.New("invalid argument"), + os.ErrNotExist, + toolPanic{value: "boom"}, + } + for _, err := range terminal { + if isRetryableToolError(err) { + t.Errorf("expected terminal (not retryable): %v", err) + } + } +} + +// TestExecutorRetrySuccessResultWithIsErrorNotRetried proves that a +// (result, nil) whose own content signals an error is NOT retried: only a +// non-nil Go error triggers retry. +func TestExecutorRetrySuccessResultWithIsErrorNotRetried(t *testing.T) { + var attempts int32 + term := false + tool := execTool{name: "toolerr", run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + atomic.AddInt32(&attempts, 1) + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("tool-level error")}, Terminate: &term}, nil + }} + cfg := newExecCfg(t, tool) + // afterToolCall marks it as an error result; still must not be retried. + isErr := true + cfg.AfterToolCall = func(ctx context.Context, call agentcore.AgentToolCall, result agentcore.AgentToolResult, isError bool) *agentcore.AfterToolCallResult { + return &agentcore.AfterToolCallResult{IsError: &isErr} + } + msg, _ := executeToolCall(context.Background(), cfg, agentcore.AgentToolCall{ID: "1", Name: "toolerr"}, nil) + if !msg.IsError { + t.Fatalf("expected error result from afterToolCall override") + } + if got := atomic.LoadInt32(&attempts); got != 1 { + t.Fatalf("(result,nil) must not be retried: got %d attempts", got) + } +} + diff --git a/pigo/internal/agenttool/tool_retry.go b/pigo/internal/agenttool/tool_retry.go new file mode 100644 index 0000000..4aaff5d --- /dev/null +++ b/pigo/internal/agenttool/tool_retry.go @@ -0,0 +1,124 @@ +// This file adds resilience to the single tool-execution seam (see +// tool_executor.go): when a tool's Execute returns a non-nil Go error, the +// error is classified as transient (worth a bounded retry) or terminal (give +// up immediately). This is deliberately separate from and does NOT touch the +// transport-layer connect-time retry in internal/provider/transport.go, which +// keeps its own, stricter semantics (only 429/503/529, respect Retry-After, +// never replay a consumed stream). +package agenttool + +import ( + "context" + "errors" + "net" + "os" + "strings" + "syscall" + "time" +) + +// maxToolRetries is the default cap on RETRIES (not attempts) for a transient +// tool error: 2 retries => at most 3 attempts total. Override per-executor via +// ToolExecutorConfig.MaxToolRetries. This is always finite; the retry loop can +// never spin forever. +const maxToolRetries = 2 + +// toolRetryBaseDelay is the unit of the small linear backoff between attempts +// (attempt N waits (N+1)*base). Kept intentionally short so retries add +// resilience without stalling the agent loop. +const toolRetryBaseDelay = 20 * time.Millisecond + +// toolPanic wraps a recovered panic value so the retry loop can (a) tell a +// panic apart from an ordinary error to shape the right message and (b) treat +// it as terminal (never retryable). +type toolPanic struct{ value any } + +func (p toolPanic) Error() string { return "panic" } + +// isRetryableToolError reports whether a non-nil error returned by a tool's +// Execute is a TRANSIENT failure worth retrying. Transient means a temporary +// IO/network/timeout condition that may succeed on a fresh attempt: +// +// - syscall.ETIMEDOUT / ECONNRESET / EAGAIN +// - a net.Error whose Timeout() or Temporary() is true +// - context.DeadlineExceeded (a per-attempt/inner deadline; the caller +// separately refuses to retry when the OUTER ctx is already done) +// - os.ErrDeadlineExceeded (i/o deadline) +// - error text containing "connection refused" / "temporarily unavailable" / +// "i/o timeout" / "connection reset" +// +// Everything else is TERMINAL and must not be retried: argument/validation +// errors, file-not-found, and — importantly — context.Canceled, which always +// means "stop", never "try again". A recovered panic (toolPanic) is terminal +// too. +func isRetryableToolError(err error) bool { + if err == nil { + return false + } + // Cancellation is always terminal, even if some inner cause looks transient. + if errors.Is(err, context.Canceled) { + return false + } + // A recovered panic is a programming error, never transient. + var tp toolPanic + if errors.As(err, &tp) { + return false + } + + // Deadline / timeout sentinels. + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, os.ErrDeadlineExceeded) { + return true + } + // Transient syscall errnos. + if errors.Is(err, syscall.ETIMEDOUT) || errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.EAGAIN) { + return true + } + // net.Error temporary/timeout conditions. + var netErr net.Error + if errors.As(err, &netErr) { + if netErr.Timeout() || netErr.Temporary() { + return true + } + } + // Best-effort string fallbacks for errors that lost their typed cause. + msg := strings.ToLower(err.Error()) + for _, s := range []string{ + "connection refused", + "temporarily unavailable", + "i/o timeout", + "connection reset", + } { + if strings.Contains(msg, s) { + return true + } + } + return false +} + +// toolRetryCap resolves the effective retry cap from the config field, +// mirroring the sentinel convention used by MaxResultBytes: 0 => default +// (maxToolRetries), <0 => disabled (0 retries, i.e. a single attempt). +func toolRetryCap(cfgMax int) int { + if cfgMax == 0 { + return maxToolRetries + } + if cfgMax < 0 { + return 0 + } + return cfgMax +} + +// waitToolRetryBackoff sleeps a small, attempt-scaled delay before the next +// attempt, but returns early (false) if ctx is cancelled/expired during the +// wait so a dead context never costs the full backoff. +func waitToolRetryBackoff(ctx context.Context, attempt int) bool { + d := time.Duration(attempt+1) * toolRetryBaseDelay + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} diff --git a/pigo/internal/agenttool/webfetch_tool.go b/pigo/internal/agenttool/webfetch_tool.go new file mode 100644 index 0000000..403329b --- /dev/null +++ b/pigo/internal/agenttool/webfetch_tool.go @@ -0,0 +1,231 @@ +// This file implements the webfetch tool (US-012, #128): fetch a URL and return +// its main text as simplified Markdown. pi has no such tool; this mirrors Claude +// Code's WebFetch. Safety properties required by the issue: +// +// - HTTP URLs are upgraded to HTTPS before the request. +// - Cross-origin redirects are NOT followed automatically; the redirect target +// is returned to the caller (the model) so it can decide whether to fetch it. +// - A request timeout and a response-body size cap bound the work. +// - A failed fetch (timeout, non-2xx, unreachable) degrades to a structured +// error result, never a panic. +// +// The optional "prompt" argument is accepted and echoed back in the result +// framing so the model keeps its intent alongside the fetched content; the tool +// does not itself call a model to summarize (that is the agent loop's job). +package agenttool + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// webFetchTimeout bounds a single fetch. webFetchMaxBytes caps how much of the +// response body is read (protects the model's context and memory from huge +// pages). webFetchMaxMarkdown caps the rendered Markdown length. +const ( + webFetchTimeout = 30 * time.Second + webFetchMaxBytes = 5 * 1024 * 1024 + webFetchMaxMarkdown = 100 * 1024 +) + +// WebFetchTool fetches a URL and returns its text as simplified Markdown. The +// zero value is usable; Client defaults to a redirect-blocking http.Client with +// webFetchTimeout. +type WebFetchTool struct { + // Client performs the HTTP request. When nil, a default client is built that + // blocks cross-origin redirects and enforces webFetchTimeout. Injected for + // tests so a fake transport can serve canned responses. + Client *http.Client +} + +// webFetchArgs is the decoded argument shape for WebFetchTool. +type webFetchArgs struct { + // URL is the page to fetch. An http:// URL is upgraded to https://. + URL string `json:"url"` + // Prompt is an optional instruction describing what the caller wants from the + // page; it is echoed into the result framing, not acted on by the tool. + Prompt string `json:"prompt,omitempty"` +} + +// Name implements AgentTool. +func (t *WebFetchTool) Name() string { return "webfetch" } + +// Description implements AgentTool. +func (t *WebFetchTool) Description() string { + return "Fetch a URL and return its main text content as simplified Markdown. " + + "HTTP URLs are upgraded to HTTPS. Cross-origin redirects are not followed; " + + "the redirect target is returned so you can fetch it explicitly. Use the " + + "optional prompt to note what you are looking for on the page." +} + +// Schema implements AgentTool. +func (t *WebFetchTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "url": {"type": "string", "description": "The URL to fetch. http:// is upgraded to https://."}, + "prompt": {"type": "string", "description": "Optional: what to extract or look for on the page."} + }, + "required": ["url"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. A fetch has no local side effects and is +// safe to run alongside other reads → parallel. +func (t *WebFetchTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} + +// errRedirectBlocked is returned by the client's CheckRedirect to stop a +// cross-origin redirect; the target is carried so Execute can report it. +type errRedirectBlocked struct{ target string } + +func (e *errRedirectBlocked) Error() string { return "cross-origin redirect blocked to " + e.target } + +// newWebFetchClient builds the default redirect-blocking client. A redirect is +// allowed only when it stays on the same host (scheme+host); a cross-origin hop +// stops with errRedirectBlocked carrying the target URL. +func newWebFetchClient() *http.Client { + return &http.Client{ + Timeout: webFetchTimeout, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + if len(via) == 0 { + return nil + } + orig := via[0].URL + if req.URL.Host != orig.Host || req.URL.Scheme != orig.Scheme { + return &errRedirectBlocked{target: req.URL.String()} + } + if len(via) >= 10 { + return errors.New("stopped after 10 redirects") + } + return nil + }, + } +} + +// Execute implements AgentTool. Fetch failures are encoded as error results (the +// returned Go error is always nil), matching the file tools' contract. +func (t *WebFetchTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[webFetchArgs](args, "webfetch") + if bad != nil { + return *bad, nil + } + raw := strings.TrimSpace(a.URL) + if raw == "" { + return errorResult("webfetch: url is required"), nil + } + + target, err := normalizeFetchURL(raw) + if err != nil { + return errorResult("webfetch: " + err.Error()), nil + } + + client := t.Client + if client == nil { + client = newWebFetchClient() + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, target, nil) + if err != nil { + return errorResult("webfetch: " + err.Error()), nil + } + req.Header.Set("User-Agent", "pigo-webfetch/1.0") + req.Header.Set("Accept", "text/html,application/xhtml+xml,text/plain;q=0.9,*/*;q=0.8") + + resp, err := client.Do(req) + if err != nil { + // A blocked cross-origin redirect is reported specially so the model can + // choose to fetch the target explicitly. errors.As unwraps the *url.Error + // http.Client wraps CheckRedirect failures in. + var blocked *errRedirectBlocked + if errors.As(err, &blocked) { + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent( + fmt.Sprintf("webfetch: cross-origin redirect not followed.\nTarget: %s\nFetch it explicitly if you want its content.", blocked.target))}, + Details: map[string]any{"redirect": blocked.target, "followed": false}, + }, nil + } + return errorResult(fmt.Sprintf("webfetch: request failed: %v", err)), nil + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return errorResult(fmt.Sprintf("webfetch: %s returned HTTP %d %s", target, resp.StatusCode, http.StatusText(resp.StatusCode))), nil + } + + body, err := io.ReadAll(io.LimitReader(resp.Body, webFetchMaxBytes)) + if err != nil { + return errorResult(fmt.Sprintf("webfetch: reading response body: %v", err)), nil + } + + ctype := resp.Header.Get("Content-Type") + var text string + if strings.Contains(ctype, "html") || looksLikeHTML(body) { + text = htmlToMarkdown(body) + } else { + text = string(body) + } + text = strings.TrimSpace(text) + truncated := false + if len(text) > webFetchMaxMarkdown { + text = text[:webFetchMaxMarkdown] + truncated = true + } + + var b strings.Builder + fmt.Fprintf(&b, "Fetched %s (HTTP %d)\n", target, resp.StatusCode) + if a.Prompt != "" { + fmt.Fprintf(&b, "Prompt: %s\n", a.Prompt) + } + if truncated { + b.WriteString("(content truncated)\n") + } + b.WriteString("\n") + b.WriteString(text) + + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(b.String())}, + Details: map[string]any{"url": target, "status": resp.StatusCode, "truncated": truncated}, + }, nil +} + +// normalizeFetchURL parses raw, upgrades an http scheme to https, and rejects +// anything that is not an absolute http(s) URL with a host. +func normalizeFetchURL(raw string) (string, error) { + u, err := url.Parse(raw) + if err != nil { + return "", fmt.Errorf("invalid url: %v", err) + } + switch u.Scheme { + case "http": + u.Scheme = "https" // upgrade + case "https": + // ok + case "": + return "", fmt.Errorf("url must be absolute with an http(s) scheme: %q", raw) + default: + return "", fmt.Errorf("unsupported url scheme %q (want http or https)", u.Scheme) + } + if u.Host == "" { + return "", fmt.Errorf("url has no host: %q", raw) + } + return u.String(), nil +} + +// looksLikeHTML sniffs whether body begins with an HTML marker, used when the +// server omits or mislabels Content-Type. +func looksLikeHTML(body []byte) bool { + head := strings.ToLower(strings.TrimSpace(string(body[:min(512, len(body))]))) + return strings.HasPrefix(head, "

hi

"), nil + }, `{"url":"http://example.com/page"}`) + if !strings.HasPrefix(gotURL, "https://") { + t.Errorf("request URL = %q, want https upgrade", gotURL) + } + if txt := agentcore.ContentToText(res.Content); !strings.Contains(txt, "hi") { + t.Errorf("result missing body text: %q", txt) + } +} + +// TestWebFetchHTMLToMarkdown checks basic HTML is reduced to Markdown. +func TestWebFetchHTMLToMarkdown(t *testing.T) { + body := `

Title

A link here.

` + res := execWebFetch(t, func(r *http.Request) (*http.Response, error) { + return makeResp(200, "text/html; charset=utf-8", body), nil + }, `{"url":"https://example.com"}`) + txt := agentcore.ContentToText(res.Content) + if !strings.Contains(txt, "# Title") { + t.Errorf("missing heading markdown in %q", txt) + } + if !strings.Contains(txt, "[link](https://x.io)") { + t.Errorf("missing link markdown in %q", txt) + } + if strings.Contains(txt, "ignore()") { + t.Errorf("script content leaked into output: %q", txt) + } +} + +// TestWebFetchNon2xxIsError checks a non-2xx status degrades to a structured +// error result (not a panic, not a Go error). +func TestWebFetchNon2xxIsError(t *testing.T) { + res := execWebFetch(t, func(r *http.Request) (*http.Response, error) { + return makeResp(404, "text/html", "not found"), nil + }, `{"url":"https://example.com/missing"}`) + txt := agentcore.ContentToText(res.Content) + if !strings.Contains(txt, "HTTP 404") { + t.Errorf("expected HTTP 404 error, got %q", txt) + } +} + +// TestWebFetchPromptEchoed checks the optional prompt is echoed into the framing. +func TestWebFetchPromptEchoed(t *testing.T) { + res := execWebFetch(t, func(r *http.Request) (*http.Response, error) { + return makeResp(200, "text/plain", "plain body"), nil + }, `{"url":"https://example.com","prompt":"find the price"}`) + if txt := agentcore.ContentToText(res.Content); !strings.Contains(txt, "Prompt: find the price") { + t.Errorf("prompt not echoed: %q", txt) + } +} + +// TestWebFetchRejectsBadScheme checks a non-http(s) scheme is rejected up front. +func TestWebFetchRejectsBadScheme(t *testing.T) { + res := execWebFetch(t, func(r *http.Request) (*http.Response, error) { + t.Fatal("transport should not be called for a bad scheme") + return nil, nil + }, `{"url":"ftp://example.com/file"}`) + if txt := agentcore.ContentToText(res.Content); !strings.Contains(txt, "unsupported url scheme") { + t.Errorf("expected scheme rejection, got %q", txt) + } +} + +// TestWebFetchMissingURL checks an empty url is rejected. +func TestWebFetchMissingURL(t *testing.T) { + tool := &WebFetchTool{} + res, err := tool.Execute(context.Background(), "c1", json.RawMessage(`{"url":" "}`), nil) + if err != nil { + t.Fatalf("Execute Go error: %v", err) + } + if txt := agentcore.ContentToText(res.Content); !strings.Contains(txt, "url is required") { + t.Errorf("expected url-required error, got %q", txt) + } +} + +// TestWebFetchCrossOriginRedirectBlocked drives the real redirect-blocking +// client (newWebFetchClient) via CheckRedirect: a cross-origin redirect must not +// be followed, and the target is reported back. +func TestWebFetchCrossOriginRedirectBlocked(t *testing.T) { + client := newWebFetchClient() + // Two hops: same-host allowed, cross-host blocked. + same := mustParse(t, "https://a.example.com/1") + cross := mustParse(t, "https://b.other.com/2") + + // Same-origin redirect: allowed (nil error). + viaSame := []*http.Request{{URL: mustParse(t, "https://a.example.com/0")}} + if err := client.CheckRedirect(&http.Request{URL: same}, viaSame); err != nil { + t.Errorf("same-origin redirect blocked unexpectedly: %v", err) + } + + // Cross-origin redirect: blocked with target carried. + viaCross := []*http.Request{{URL: mustParse(t, "https://a.example.com/0")}} + err := client.CheckRedirect(&http.Request{URL: cross}, viaCross) + var blocked *errRedirectBlocked + if err == nil || !errorAsRedirect(err, &blocked) { + t.Fatalf("cross-origin redirect not blocked: %v", err) + } + if blocked.target != "https://b.other.com/2" { + t.Errorf("blocked target = %q", blocked.target) + } +} + +// TestNormalizeFetchURL covers the scheme/host rules directly. +func TestNormalizeFetchURL(t *testing.T) { + cases := []struct { + in, want string + wantErr bool + }{ + {"http://x.com/a", "https://x.com/a", false}, + {"https://x.com", "https://x.com", false}, + {"x.com/a", "", true}, // no scheme + {"ftp://x.com", "", true}, // bad scheme + {"https://", "", true}, // no host + } + for _, c := range cases { + got, err := normalizeFetchURL(c.in) + if c.wantErr { + if err == nil { + t.Errorf("normalizeFetchURL(%q) = %q, want error", c.in, got) + } + continue + } + if err != nil { + t.Errorf("normalizeFetchURL(%q) error: %v", c.in, err) + continue + } + if got != c.want { + t.Errorf("normalizeFetchURL(%q) = %q, want %q", c.in, got, c.want) + } + } +} diff --git a/pigo/internal/agenttool/websearch_backends.go b/pigo/internal/agenttool/websearch_backends.go new file mode 100644 index 0000000..3312fa5 --- /dev/null +++ b/pigo/internal/agenttool/websearch_backends.go @@ -0,0 +1,275 @@ +// This file implements the websearch backends: Tavily and Brave (credentialed +// JSON APIs) plus a keyless DuckDuckGo HTML fallback. selectSearchBackend picks +// the first backend whose credential is present, defaulting to DuckDuckGo. +package agenttool + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + "golang.org/x/net/html" +) + +// searchBackend is one pluggable search provider. name is used in result framing +// and error messages; search runs the query and returns up to count normalized +// hits. +type searchBackend interface { + name() string + search(ctx context.Context, client *http.Client, query string, count int) ([]searchResult, error) +} + +// selectSearchBackend returns the first backend whose credential env var is set, +// falling back to the keyless DuckDuckGo backend. The order encodes preference: +// LLM-optimized Tavily first, then Brave, then the keyless fallback. +func selectSearchBackend(getenv func(string) string) searchBackend { + if k := strings.TrimSpace(getenv("TAVILY_API_KEY")); k != "" { + return tavilyBackend{apiKey: k} + } + if k := strings.TrimSpace(getenv("BRAVE_API_KEY")); k != "" { + return braveBackend{apiKey: k} + } + return duckDuckGoBackend{} +} + +// searchBodyLimit caps how much of a backend response body is read. +const searchBodyLimit = 4 * 1024 * 1024 + +// --- Tavily --------------------------------------------------------------- + +type tavilyBackend struct{ apiKey string } + +func (b tavilyBackend) name() string { return "tavily" } + +func (b tavilyBackend) search(ctx context.Context, client *http.Client, query string, count int) ([]searchResult, error) { + reqBody, _ := json.Marshal(map[string]any{"query": query, "max_results": count}) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://api.tavily.com/search", bytes.NewReader(reqBody)) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+b.apiKey) + req.Header.Set("Content-Type", "application/json") + + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, searchBodyLimit)) + if err != nil { + return nil, err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + + var decoded struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` + } `json:"results"` + } + if err := json.Unmarshal(body, &decoded); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + out := make([]searchResult, 0, len(decoded.Results)) + for _, r := range decoded.Results { + out = append(out, searchResult{Title: r.Title, URL: r.URL, Snippet: r.Content}) + } + return out, nil +} + +// --- Brave ---------------------------------------------------------------- + +type braveBackend struct{ apiKey string } + +func (b braveBackend) name() string { return "brave" } + +func (b braveBackend) search(ctx context.Context, client *http.Client, query string, count int) ([]searchResult, error) { + u := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", url.QueryEscape(query), count) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/json") + req.Header.Set("X-Subscription-Token", b.apiKey) + + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, searchBodyLimit)) + if err != nil { + return nil, err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + + var decoded struct { + Web struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Description string `json:"description"` + } `json:"results"` + } `json:"web"` + } + if err := json.Unmarshal(body, &decoded); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + out := make([]searchResult, 0, len(decoded.Web.Results)) + for _, r := range decoded.Web.Results { + out = append(out, searchResult{Title: stripHTMLTags(r.Title), URL: r.URL, Snippet: stripHTMLTags(r.Description)}) + } + return out, nil +} + +// --- DuckDuckGo (keyless fallback) ---------------------------------------- + +type duckDuckGoBackend struct{} + +func (duckDuckGoBackend) name() string { return "duckduckgo" } + +func (duckDuckGoBackend) search(ctx context.Context, client *http.Client, query string, count int) ([]searchResult, error) { + u := "https://html.duckduckgo.com/html/?q=" + url.QueryEscape(query) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil) + if err != nil { + return nil, err + } + // A browser-like User-Agent avoids the endpoint serving an empty/blocked page. + req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; pigo-websearch/1.0)") + req.Header.Set("Accept", "text/html") + + resp, err := client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, searchBodyLimit)) + if err != nil { + return nil, err + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("HTTP %d", resp.StatusCode) + } + + results, err := parseDuckDuckGoHTML(body) + if err != nil { + return nil, err + } + if len(results) > count { + results = results[:count] + } + return results, nil +} + +// parseDuckDuckGoHTML extracts result links and snippets from the DuckDuckGo +// HTML endpoint. Title/URL come from ; the URL is wrapped in +// a redirect carrying the real target in the uddg query param, which is decoded. +// Snippets come from elements with class "result__snippet", matched to results by +// position. A result with no snippet is still returned (snippet empty). +func parseDuckDuckGoHTML(body []byte) ([]searchResult, error) { + doc, err := html.Parse(bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("parsing html: %w", err) + } + var results []searchResult + var snippets []string + var walk func(*html.Node) + walk = func(n *html.Node) { + if n.Type == html.ElementNode && n.Data == "a" && hasClass(n, "result__a") { + href := attr(n, "href") + results = append(results, searchResult{Title: nodeText(n), URL: unwrapDDGURL(href)}) + } + if n.Type == html.ElementNode && hasClass(n, "result__snippet") { + snippets = append(snippets, nodeText(n)) + } + for c := n.FirstChild; c != nil; c = c.NextSibling { + walk(c) + } + } + walk(doc) + for i := range results { + if i < len(snippets) { + results[i].Snippet = snippets[i] + } + } + return results, nil +} + +// unwrapDDGURL turns a DuckDuckGo redirect href ("//duckduckgo.com/l/?uddg=...") +// into the real target by decoding the uddg param. A non-redirect href is +// returned as-is (with a scheme added when protocol-relative). +func unwrapDDGURL(href string) string { + raw := href + if strings.HasPrefix(raw, "//") { + raw = "https:" + raw + } + u, err := url.Parse(raw) + if err != nil { + return href + } + if target := u.Query().Get("uddg"); target != "" { + return target + } + return raw +} + +// --- HTML helpers --------------------------------------------------------- + +// attr returns the value of the named attribute on n, or "". +func attr(n *html.Node, name string) string { + for _, a := range n.Attr { + if a.Key == name { + return a.Val + } + } + return "" +} + +// hasClass reports whether n's class attribute contains the given class token. +func hasClass(n *html.Node, class string) bool { + for _, f := range strings.Fields(attr(n, "class")) { + if f == class { + return true + } + } + return false +} + +// nodeText returns the concatenated, space-collapsed text content of n. +func nodeText(n *html.Node) string { + var b strings.Builder + var walk func(*html.Node) + walk = func(x *html.Node) { + if x.Type == html.TextNode { + b.WriteString(x.Data) + } + for c := x.FirstChild; c != nil; c = c.NextSibling { + walk(c) + } + } + walk(n) + return strings.Join(strings.Fields(b.String()), " ") +} + +// stripHTMLTags removes inline markup (e.g. Brave's highlights) from a +// snippet, leaving space-collapsed text. Malformed fragments are returned as-is. +func stripHTMLTags(s string) string { + if !strings.ContainsRune(s, '<') { + return s + } + doc, err := html.Parse(strings.NewReader(s)) + if err != nil { + return s + } + return nodeText(doc) +} diff --git a/pigo/internal/agenttool/websearch_tool.go b/pigo/internal/agenttool/websearch_tool.go new file mode 100644 index 0000000..1795d32 --- /dev/null +++ b/pigo/internal/agenttool/websearch_tool.go @@ -0,0 +1,217 @@ +// This file implements the websearch tool: run a web search and return the top +// results (title, URL, snippet) as Markdown. pi has no such tool; this mirrors +// Claude Code's WebSearch. It is provider-agnostic and auto-detects a backend by +// available credentials so it works out of the box: +// +// - Tavily when TAVILY_API_KEY is set (LLM-optimized results). +// - Brave when BRAVE_API_KEY is set (independent index). +// - DuckDuckGo as a keyless fallback (HTML endpoint, no API key needed). +// +// The first backend whose credential is present wins; DuckDuckGo is always the +// last-resort fallback. Optional allowed/blocked domain filters are applied +// uniformly to every backend by post-filtering the result hosts, so behavior is +// consistent regardless of which backend served the query. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "os" + "strings" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// webSearchTimeout bounds a single search request. webSearchDefaultCount is used +// when the caller omits count; webSearchMaxCount caps it so a run cannot pull an +// unbounded result set into the model's context. +const ( + webSearchTimeout = 15 * time.Second + webSearchDefaultCount = 5 + webSearchMaxCount = 10 +) + +// WebSearchTool runs a web search via the first available backend. The zero +// value is usable: Client defaults to an http.Client with webSearchTimeout and +// getenv defaults to os.Getenv (both injected in tests). +type WebSearchTool struct { + // Client performs backend HTTP requests. When nil, a default client bounded by + // webSearchTimeout is built. Injected for tests to serve canned responses. + Client *http.Client + // getenv reads credentials for backend selection. When nil, os.Getenv is used. + // Injected for tests so backend selection is deterministic without touching the + // process environment. + getenv func(string) string +} + +// webSearchArgs is the decoded argument shape for WebSearchTool. +type webSearchArgs struct { + // Query is the search query (required). + Query string `json:"query"` + // Count is the desired number of results (optional; clamped to webSearchMaxCount). + Count int `json:"count,omitempty"` + // AllowedDomains, when non-empty, keeps only results whose host matches one of + // these domains (suffix match). BlockedDomains drops results whose host matches. + AllowedDomains []string `json:"allowed_domains,omitempty"` + BlockedDomains []string `json:"blocked_domains,omitempty"` +} + +// searchResult is one normalized hit shared across backends. +type searchResult struct { + Title string + URL string + Snippet string +} + +// Name implements AgentTool. +func (t *WebSearchTool) Name() string { return "websearch" } + +// Description implements AgentTool. +func (t *WebSearchTool) Description() string { + return "Search the web and return the top results (title, URL, snippet). " + + "Auto-selects a backend by available credentials (Tavily, Brave, or a " + + "keyless DuckDuckGo fallback). Use allowed_domains/blocked_domains to " + + "restrict results by host. Follow up with the webfetch tool to read a result." +} + +// Schema implements AgentTool. +func (t *WebSearchTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "query": {"type": "string", "description": "The search query."}, + "count": {"type": "integer", "description": "Desired number of results (default 5, max 10).", "minimum": 1, "maximum": 10}, + "allowed_domains": {"type": "array", "items": {"type": "string"}, "description": "Only include results from these domains (suffix match)."}, + "blocked_domains": {"type": "array", "items": {"type": "string"}, "description": "Exclude results from these domains (suffix match)."} + }, + "required": ["query"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. A search has no local side effects and is +// safe to run alongside other reads → parallel. +func (t *WebSearchTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} + +// Execute implements AgentTool. Backend failures are encoded as error results +// (the returned Go error is always nil), matching the file tools' contract. +func (t *WebSearchTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[webSearchArgs](args, "websearch") + if bad != nil { + return *bad, nil + } + query := strings.TrimSpace(a.Query) + if query == "" { + return errorResult("websearch: query is required"), nil + } + + count := a.Count + if count <= 0 { + count = webSearchDefaultCount + } + if count > webSearchMaxCount { + count = webSearchMaxCount + } + + client := t.Client + if client == nil { + client = &http.Client{Timeout: webSearchTimeout} + } + getenv := t.getenv + if getenv == nil { + getenv = os.Getenv + } + + backend := selectSearchBackend(getenv) + // A domain-filtered query can discard most raw hits, so over-fetch before + // filtering to still land near the requested count. + fetchCount := count + if len(a.AllowedDomains) > 0 || len(a.BlockedDomains) > 0 { + fetchCount = min(webSearchMaxCount, count*3) + } + + results, err := backend.search(ctx, client, query, fetchCount) + if err != nil { + return errorResult(fmt.Sprintf("websearch: %s backend failed: %v", backend.name(), err)), nil + } + results = filterByDomain(results, a.AllowedDomains, a.BlockedDomains) + if len(results) > count { + results = results[:count] + } + + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(renderSearchResults(query, backend.name(), results))}, + Details: map[string]any{"backend": backend.name(), "query": query, "count": len(results)}, + }, nil +} + +// renderSearchResults formats the hits as a numbered Markdown list, noting which +// backend served the query so the model knows the source. +func renderSearchResults(query, backend string, results []searchResult) string { + var b strings.Builder + fmt.Fprintf(&b, "Search results for %q (via %s):\n", query, backend) + if len(results) == 0 { + b.WriteString("\n(no results)") + return b.String() + } + for i, r := range results { + fmt.Fprintf(&b, "\n%d. %s\n %s\n", i+1, strings.TrimSpace(r.Title), strings.TrimSpace(r.URL)) + if s := strings.TrimSpace(r.Snippet); s != "" { + fmt.Fprintf(&b, " %s\n", s) + } + } + return strings.TrimRight(b.String(), "\n") +} + +// filterByDomain keeps only results whose host suffix-matches an allowed domain +// (when allowed is non-empty) and drops any whose host suffix-matches a blocked +// domain. An unparseable URL is dropped only under an allow-list. +func filterByDomain(results []searchResult, allowed, blocked []string) []searchResult { + if len(allowed) == 0 && len(blocked) == 0 { + return results + } + out := results[:0:0] + for _, r := range results { + host := hostOf(r.URL) + if len(allowed) > 0 && !matchesAnyDomain(host, allowed) { + continue + } + if matchesAnyDomain(host, blocked) { + continue + } + out = append(out, r) + } + return out +} + +// hostOf extracts the lowercased host from a result URL, or "" if unparseable. +func hostOf(raw string) string { + u, err := url.Parse(raw) + if err != nil { + return "" + } + return strings.ToLower(u.Hostname()) +} + +// matchesAnyDomain reports whether host equals or is a subdomain of any domain. +func matchesAnyDomain(host string, domains []string) bool { + if host == "" { + return false + } + for _, d := range domains { + d = strings.ToLower(strings.TrimSpace(d)) + if d == "" { + continue + } + if host == d || strings.HasSuffix(host, "."+d) { + return true + } + } + return false +} diff --git a/pigo/internal/agenttool/websearch_tool_test.go b/pigo/internal/agenttool/websearch_tool_test.go new file mode 100644 index 0000000..8cd5058 --- /dev/null +++ b/pigo/internal/agenttool/websearch_tool_test.go @@ -0,0 +1,160 @@ +// Tests for the websearch tool: backend auto-selection by credential, per-backend +// response parsing (Tavily JSON, Brave JSON with HTML highlights, DuckDuckGo HTML +// with redirect-wrapped URLs), domain filtering, count clamping, and structured +// errors. A fake RoundTripper serves canned responses so no network is touched. +package agenttool + +import ( + "context" + "encoding/json" + "net/http" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// execWebSearch runs the tool with a fake transport and a fixed environment. +func execWebSearch(t *testing.T, env map[string]string, fn roundTripFunc, args string) agentcore.AgentToolResult { + t.Helper() + tool := &WebSearchTool{ + Client: &http.Client{Transport: fn}, + getenv: func(k string) string { return env[k] }, + } + res, err := tool.Execute(context.Background(), "c1", json.RawMessage(args), nil) + if err != nil { + t.Fatalf("Execute returned Go error: %v", err) + } + return res +} + +// resultText is defined in read_tool_test.go (same package). + +func TestSelectSearchBackend(t *testing.T) { + cases := []struct { + env map[string]string + want string + }{ + {map[string]string{"TAVILY_API_KEY": "t"}, "tavily"}, + {map[string]string{"BRAVE_API_KEY": "b"}, "brave"}, + {map[string]string{"TAVILY_API_KEY": "t", "BRAVE_API_KEY": "b"}, "tavily"}, + {map[string]string{}, "duckduckgo"}, + } + for _, c := range cases { + got := selectSearchBackend(func(k string) string { return c.env[k] }).name() + if got != c.want { + t.Errorf("env %v: backend = %q, want %q", c.env, got, c.want) + } + } +} + +func TestWebSearchTavily(t *testing.T) { + body := `{"results":[{"title":"Go","url":"https://go.dev","content":"The Go language"},{"title":"Docs","url":"https://pkg.go.dev","content":"packages"}]}` + var gotAuth string + res := execWebSearch(t, map[string]string{"TAVILY_API_KEY": "secret"}, + func(r *http.Request) (*http.Response, error) { + if r.URL.Host != "api.tavily.com" { + t.Errorf("unexpected host %q", r.URL.Host) + } + gotAuth = r.Header.Get("Authorization") + return makeResp(200, "application/json", body), nil + }, `{"query":"go language"}`) + + txt := resultText(res) + if gotAuth != "Bearer secret" { + t.Errorf("Authorization = %q, want Bearer secret", gotAuth) + } + if !strings.Contains(txt, "via tavily") || !strings.Contains(txt, "https://go.dev") || !strings.Contains(txt, "The Go language") { + t.Errorf("unexpected result:\n%s", txt) + } + if bk, _ := res.Details.(map[string]any)["backend"].(string); bk != "tavily" { + t.Errorf("Details.backend = %q, want tavily", bk) + } +} + +func TestWebSearchBraveStripsHTML(t *testing.T) { + body := `{"web":{"results":[{"title":"Rust lang","url":"https://rust-lang.org","description":"A systems language"}]}}` + var gotToken string + res := execWebSearch(t, map[string]string{"BRAVE_API_KEY": "tok"}, + func(r *http.Request) (*http.Response, error) { + if r.URL.Host != "api.search.brave.com" { + t.Errorf("unexpected host %q", r.URL.Host) + } + gotToken = r.Header.Get("X-Subscription-Token") + return makeResp(200, "application/json", body), nil + }, `{"query":"rust"}`) + + txt := resultText(res) + if gotToken != "tok" { + t.Errorf("X-Subscription-Token = %q, want tok", gotToken) + } + if strings.Contains(txt, "") { + t.Errorf("HTML tags not stripped:\n%s", txt) + } + if !strings.Contains(txt, "Rust lang") || !strings.Contains(txt, "A systems language") { + t.Errorf("unexpected result:\n%s", txt) + } +} + +func TestWebSearchDuckDuckGo(t *testing.T) { + html := ` + ` + res := execWebSearch(t, map[string]string{}, + func(r *http.Request) (*http.Response, error) { + if r.URL.Host != "html.duckduckgo.com" { + t.Errorf("unexpected host %q", r.URL.Host) + } + return makeResp(200, "text/html", html), nil + }, `{"query":"anything"}`) + + txt := resultText(res) + if !strings.Contains(txt, "https://example.com/a") || !strings.Contains(txt, "https://example.org/b") { + t.Errorf("redirect URLs not decoded:\n%s", txt) + } + if !strings.Contains(txt, "First Title") || !strings.Contains(txt, "Second snippet") { + t.Errorf("titles/snippets missing:\n%s", txt) + } +} + +func TestWebSearchDomainFilter(t *testing.T) { + body := `{"results":[{"title":"A","url":"https://keep.com/x","content":"a"},{"title":"B","url":"https://drop.com/y","content":"b"}]}` + res := execWebSearch(t, map[string]string{"TAVILY_API_KEY": "k"}, + func(r *http.Request) (*http.Response, error) { + return makeResp(200, "application/json", body), nil + }, `{"query":"q","allowed_domains":["keep.com"]}`) + + txt := resultText(res) + if strings.Contains(txt, "drop.com") { + t.Errorf("blocked domain leaked:\n%s", txt) + } + if !strings.Contains(txt, "keep.com") { + t.Errorf("allowed domain dropped:\n%s", txt) + } +} + +func TestWebSearchEmptyQuery(t *testing.T) { + res := execWebSearch(t, map[string]string{}, + func(r *http.Request) (*http.Response, error) { + t.Error("transport should not be called for empty query") + return makeResp(200, "text/html", ""), nil + }, `{"query":" "}`) + if !strings.Contains(resultText(res), "query is required") { + t.Errorf("want query-required error, got:\n%s", resultText(res)) + } +} + +func TestWebSearchBackendError(t *testing.T) { + res := execWebSearch(t, map[string]string{"TAVILY_API_KEY": "k"}, + func(r *http.Request) (*http.Response, error) { + return makeResp(500, "application/json", "boom"), nil + }, `{"query":"q"}`) + if !strings.Contains(resultText(res), "tavily backend failed") { + t.Errorf("want backend-failed error, got:\n%s", resultText(res)) + } +} diff --git a/pigo/internal/agenttool/write_tool.go b/pigo/internal/agenttool/write_tool.go new file mode 100644 index 0000000..ba9557e --- /dev/null +++ b/pigo/internal/agenttool/write_tool.go @@ -0,0 +1,125 @@ +// This file implements the write tool (US-016): create or overwrite a file at a +// given path, creating parent directories as needed. Overwrites are reported so +// the caller/model knows an existing file was replaced (parity with pi's write +// behavior). Paths resolve against a Root and are rejected if they escape it. +package agenttool + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// WriteTool writes text files under Root, creating parent directories as needed. +type WriteTool struct { + // Root bounds all writes; a path resolving outside Root is rejected. Empty + // Root defaults to the current working directory. + Root string + // ExtraRoots are additional trusted directories a write may target even though + // they lie outside Root. It exists for the skills directory so the model can + // author or update skills (create a new SKILL.md, edit an existing one) that + // live outside the workspace. + ExtraRoots []string + // Snap, when non-nil, records the file's prior content before it is written so + // the /rewind command can roll the change back. It is shared with the edit tool. + Snap *FileSnapshotRecorder +} + +// writeToolArgs is the decoded argument shape for WriteTool. +type writeToolArgs struct { + // Path is the file to write, relative to Root (or absolute within Root). + Path string `json:"path"` + // Content is the full file contents to write (overwrites any existing file). + Content string `json:"content"` +} + +// Name implements AgentTool. +func (t *WriteTool) Name() string { return "write" } + +// Description implements AgentTool. +func (t *WriteTool) Description() string { + return "Create or overwrite a file at the given path, creating parent " + + "directories as needed. Overwriting an existing file is reported." +} + +// Schema implements AgentTool. +func (t *WriteTool) Schema() json.RawMessage { + return json.RawMessage(`{ + "type": "object", + "properties": { + "path": {"type": "string", "description": "File path to write, relative to the workspace root."}, + "content": {"type": "string", "description": "Full file contents to write."} + }, + "required": ["path", "content"], + "additionalProperties": false +}`) +} + +// ExecutionMode implements AgentTool. Writes mutate the filesystem → sequential +// so a batch does not race concurrent writes to the same tree. +func (t *WriteTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// resolvePath resolves p against Root (or any ExtraRoots) via the shared +// resolveWithin boundary policy, so every file tool enforces the same +// workspace-escape guard while writes can also reach trusted extra roots. +func (t *WriteTool) resolvePath(p string) (string, error) { + if len(t.ExtraRoots) == 0 { + return resolveWithin(t.Root, p) + } + return resolveWithinAny(append([]string{t.Root}, t.ExtraRoots...), p) +} + +// Execute implements AgentTool. Write failures are encoded as error results; +// the returned Go error is reserved for nothing here (argument decode also +// degrades to a result), matching the read tool's contract. +func (t *WriteTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + a, bad := decodeArgs[writeToolArgs](args, "write") + if bad != nil { + return *bad, nil + } + if a.Path == "" { + return errorResult("write: path is required"), nil + } + full, err := t.resolvePath(a.Path) + if err != nil { + return errorResult("write: " + err.Error()), nil + } + + // Detect overwrite before writing so the result can report it. A path that + // points at a directory is an error, not an overwrite. + overwrote := false + if info, statErr := os.Stat(full); statErr == nil { + if info.IsDir() { + return errorResult(fmt.Sprintf("write: %q is a directory, not a file", a.Path)), nil + } + overwrote = true + } + + // Create parent directories as needed. + if dir := filepath.Dir(full); dir != "" { + if err := os.MkdirAll(dir, dirPerm); err != nil { + return errorResult(fmt.Sprintf("write: cannot create parent directories for %q: %v", a.Path, err)), nil + } + } + + // Snapshot the prior state before mutating so /rewind can restore it. + t.Snap.Record(full) + if err := os.WriteFile(full, []byte(a.Content), filePerm); err != nil { + return errorResult(fmt.Sprintf("write: cannot write %q: %v", a.Path, err)), nil + } + verb := "Created" + if overwrote { + verb = "Overwrote" + } + msg := fmt.Sprintf("%s %s (%d bytes)", verb, a.Path, len(a.Content)) + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(msg)}, + Details: map[string]any{"path": a.Path, "bytes": len(a.Content), "overwrote": overwrote}, + }, nil +} diff --git a/pigo/internal/agenttool/write_tool_test.go b/pigo/internal/agenttool/write_tool_test.go new file mode 100644 index 0000000..b7276f9 --- /dev/null +++ b/pigo/internal/agenttool/write_tool_test.go @@ -0,0 +1,164 @@ +package agenttool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func runWrite(t *testing.T, tool *WriteTool, args map[string]any) agentcore.AgentToolResult { + t.Helper() + raw, err := json.Marshal(args) + if err != nil { + t.Fatalf("marshal args: %v", err) + } + res, gerr := tool.Execute(context.Background(), "call-1", raw, nil) + if gerr != nil { + t.Fatalf("execute returned go error: %v", gerr) + } + return res +} + +func TestWriteToolCreate(t *testing.T) { + dir := t.TempDir() + tool := &WriteTool{Root: dir} + res := runWrite(t, tool, map[string]any{"path": "out.txt", "content": "hello"}) + if !strings.Contains(resultText(res), "Created") { + t.Errorf("expected Created, got %q", resultText(res)) + } + got, err := os.ReadFile(filepath.Join(dir, "out.txt")) + if err != nil || string(got) != "hello" { + t.Errorf("file content = %q, err = %v", got, err) + } +} + +func TestWriteToolCreatesParentDirs(t *testing.T) { + dir := t.TempDir() + tool := &WriteTool{Root: dir} + res := runWrite(t, tool, map[string]any{"path": "a/b/c/deep.txt", "content": "x"}) + if strings.Contains(resultText(res), "error") { + t.Errorf("unexpected error: %q", resultText(res)) + } + if _, err := os.Stat(filepath.Join(dir, "a", "b", "c", "deep.txt")); err != nil { + t.Errorf("nested file not created: %v", err) + } +} + +func TestWriteToolOverwrite(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "exists.txt") + if err := os.WriteFile(path, []byte("old"), 0o644); err != nil { + t.Fatalf("seed: %v", err) + } + tool := &WriteTool{Root: dir} + res := runWrite(t, tool, map[string]any{"path": "exists.txt", "content": "new"}) + if !strings.Contains(resultText(res), "Overwrote") { + t.Errorf("expected Overwrote, got %q", resultText(res)) + } + got, _ := os.ReadFile(path) + if string(got) != "new" { + t.Errorf("content = %q, want new", got) + } + // Details should flag the overwrite. + details, ok := res.Details.(map[string]any) + if !ok || details["overwrote"] != true { + t.Errorf("details missing overwrote flag: %+v", res.Details) + } +} + +func TestWriteToolPathTraversal(t *testing.T) { + dir := t.TempDir() + tool := &WriteTool{Root: dir} + res := runWrite(t, tool, map[string]any{"path": "../escape.txt", "content": "x"}) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Errorf("expected boundary error, got %q", resultText(res)) + } + // The escape file must not exist. + if _, err := os.Stat(filepath.Join(filepath.Dir(dir), "escape.txt")); err == nil { + t.Fatal("path traversal wrote outside the root!") + } +} + +func TestWriteToolExtraRootsAllowsSkillAuthoring(t *testing.T) { + work := t.TempDir() + skills := t.TempDir() + + // Without ExtraRoots, authoring a skill outside the workspace is rejected. + target := filepath.Join(skills, "newskill", "SKILL.md") + bounded := &WriteTool{Root: work} + res := runWrite(t, bounded, map[string]any{"path": target, "content": "x"}) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Fatalf("expected boundary rejection without ExtraRoots, got %q", resultText(res)) + } + if _, err := os.Stat(target); err == nil { + t.Fatal("write escaped the workspace without ExtraRoots!") + } + + // With the skills dir as an extra root, the new skill file (and its parent + // dirs) is created. + tool := &WriteTool{Root: work, ExtraRoots: []string{skills}} + res = runWrite(t, tool, map[string]any{"path": target, "content": "skill body"}) + if strings.Contains(resultText(res), "error") { + t.Fatalf("unexpected error authoring skill: %q", resultText(res)) + } + got, err := os.ReadFile(target) + if err != nil || string(got) != "skill body" { + t.Fatalf("skill file content = %q, err = %v", got, err) + } +} + +func TestWriteToolExtraRootsStillBlocksUntrustedPath(t *testing.T) { + work := t.TempDir() + skills := t.TempDir() + other := t.TempDir() + target := filepath.Join(other, "escape.txt") + + tool := &WriteTool{Root: work, ExtraRoots: []string{skills}} + res := runWrite(t, tool, map[string]any{"path": target, "content": "x"}) + if !strings.Contains(resultText(res), "outside the workspace root") { + t.Errorf("expected boundary error for untrusted path, got %q", resultText(res)) + } + if _, err := os.Stat(target); err == nil { + t.Fatal("write escaped both roots!") + } +} + +func TestWriteToolDirectoryTarget(t *testing.T) { + dir := t.TempDir() + sub := filepath.Join(dir, "adir") + if err := os.Mkdir(sub, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + tool := &WriteTool{Root: dir} + res := runWrite(t, tool, map[string]any{"path": "adir", "content": "x"}) + if !strings.Contains(resultText(res), "is a directory") { + t.Errorf("expected directory error, got %q", resultText(res)) + } +} + +func TestWriteToolMissingArgs(t *testing.T) { + tool := &WriteTool{Root: t.TempDir()} + res := runWrite(t, tool, map[string]any{"content": "x"}) + if !strings.Contains(resultText(res), "path is required") { + t.Errorf("expected path-required error, got %q", resultText(res)) + } +} + +func TestWriteToolMode(t *testing.T) { + tool := &WriteTool{} + if tool.Name() != "write" { + t.Errorf("name = %q", tool.Name()) + } + if tool.ExecutionMode() != agentcore.ToolExecutionSequential { + t.Error("write should be sequential") + } + var schema map[string]any + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Errorf("schema not valid JSON: %v", err) + } +} diff --git a/pigo/internal/builtinskills/bootstrap.go b/pigo/internal/builtinskills/bootstrap.go new file mode 100644 index 0000000..10290fd --- /dev/null +++ b/pigo/internal/builtinskills/bootstrap.go @@ -0,0 +1,177 @@ +// This file holds the first-run bootstrap: on the first launch (per pigo home), +// the built-in skill collections in Manifest are copied into the user's skills +// directory so they load as /skill-name commands with no manual install. +// +// The flow is designed to be silent and non-blocking (a failed bootstrap must +// never stop pigo from starting) and idempotent (skills already present are left +// untouched, so a user's edits are never clobbered). A state file under the pigo +// home records which collections+versions have been installed, so a completed +// bootstrap is skipped on later launches and a bumped collection Version +// re-triggers installation of any still-missing skills. +package builtinskills + +import ( + "encoding/json" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" +) + +// stateFileName is the bootstrap record kept under the pigo home directory. It +// maps a collection name to the version last installed, so Bootstrap can decide +// per collection whether work remains. +const stateFileName = "builtin-skills.json" + +// state is the on-disk bootstrap record: collection name -> installed version. +type state struct { + Installed map[string]string `json:"installed"` +} + +// Bootstrap installs any not-yet-installed built-in skill collections into +// skillsDir, recording progress under pigoHome. It is safe to call on every +// launch: collections already recorded at their current Version are skipped, and +// individual skills whose target directory already exists are left untouched. +// +// It never returns an error — bootstrap is best-effort and must not block +// startup — but writes a one-line note per failure to logw when logw is non-nil +// (callers pass a writer only in debug/verbose mode, keeping normal runs silent). +// Empty pigoHome or skillsDir disables bootstrap (home unresolved): nothing is +// installed and nothing is logged. +func Bootstrap(pigoHome, skillsDir string, logw io.Writer) { + bootstrap(Manifest(), pigoHome, skillsDir, logw) +} + +// bootstrap is the testable core of Bootstrap, parameterized on the manifest so +// tests can inject synthetic collections without touching the embedded set. +func bootstrap(sets []Set, pigoHome, skillsDir string, logw io.Writer) { + logf := func(format string, a ...any) { + if logw != nil { + fmt.Fprintf(logw, format, a...) + } + } + if pigoHome == "" || skillsDir == "" { + return // home unresolved; nothing we can safely do + } + + st := loadState(filepath.Join(pigoHome, stateFileName)) + + changed := false + for _, set := range sets { + // Per-collection version gate: once recorded at this Version the whole + // set is skipped, so a skill the user later *deletes* is not restored + // (only a Version bump re-triggers install of still-missing skills). + // This is deliberate — silently re-adding a removed skill would fight + // the user's choice. An empty Version is never "already installed" + // (the zero-value lookup would otherwise equal it and skip forever), + // so a set with a blank Version always attempts install. + if set.Version != "" && st.Installed[set.Name] == set.Version { + continue + } + // installSet reports whether the collection is fully installed (every + // skill now present on disk). Only then do we record the version, so a + // partial failure re-runs next launch (satisfying "retry on next run"). + if installSet(set, skillsDir, logf) { + if st.Installed == nil { + st.Installed = map[string]string{} + } + st.Installed[set.Name] = set.Version + changed = true + } + } + + if changed { + if err := saveState(pigoHome, st); err != nil { + logf("builtinskills: could not save state: %v\n", err) + } + } +} + +// installSet copies each skill in the set into skillsDir//, skipping any +// whose target directory already exists (never clobbering a user's copy). A +// single skill's failure does not abort the rest. It returns true only when +// every named skill is present on disk afterward, so the caller can decide +// whether to record the collection as fully installed. +func installSet(set Set, skillsDir string, logf func(string, ...any)) bool { + if err := os.MkdirAll(skillsDir, 0o755); err != nil { + logf("builtinskills: create skills dir %q: %v\n", skillsDir, err) + return false + } + allPresent := true + for _, name := range set.Skills { + dest := filepath.Join(skillsDir, name) + if _, err := os.Stat(dest); err == nil { + continue // already present (user copy or prior install): leave it + } + src := path.Join(set.Root, name) + if err := copyTree(set.FS, src, dest); err != nil { + logf("builtinskills: install %q: %v\n", name, err) + // Remove a half-written tree so a later run starts clean. + _ = os.RemoveAll(dest) + allPresent = false + continue + } + } + return allPresent +} + +// copyTree copies the directory tree rooted at src within srcFS to the on-disk +// directory dest, creating parents as needed. Files are written 0o644 and +// directories 0o755. +func copyTree(srcFS fs.FS, src, dest string) error { + entries, err := fs.ReadDir(srcFS, src) + if err != nil { + return err + } + if err := os.MkdirAll(dest, 0o755); err != nil { + return err + } + for _, e := range entries { + srcPath := path.Join(src, e.Name()) + destPath := filepath.Join(dest, e.Name()) + if e.IsDir() { + if err := copyTree(srcFS, srcPath, destPath); err != nil { + return err + } + continue + } + data, err := fs.ReadFile(srcFS, srcPath) + if err != nil { + return err + } + if err := os.WriteFile(destPath, data, 0o644); err != nil { + return err + } + } + return nil +} + +// loadState reads the bootstrap record at p. A missing or malformed file yields +// an empty state (treated as "nothing installed"), so a corrupt record just +// re-triggers a — idempotent — install rather than failing. +func loadState(p string) state { + data, err := os.ReadFile(p) + if err != nil { + return state{} + } + var st state + if json.Unmarshal(data, &st) != nil { + return state{} + } + return st +} + +// saveState writes the bootstrap record under pigoHome, creating the home +// directory if needed. +func saveState(pigoHome string, st state) error { + if err := os.MkdirAll(pigoHome, 0o755); err != nil { + return err + } + data, err := json.MarshalIndent(st, "", " ") + if err != nil { + return err + } + return os.WriteFile(filepath.Join(pigoHome, stateFileName), data, 0o644) +} diff --git a/pigo/internal/builtinskills/bootstrap_test.go b/pigo/internal/builtinskills/bootstrap_test.go new file mode 100644 index 0000000..1c811ea --- /dev/null +++ b/pigo/internal/builtinskills/bootstrap_test.go @@ -0,0 +1,164 @@ +package builtinskills + +import ( + "os" + "path" + "path/filepath" + "testing" + "testing/fstest" +) + +// mapSet builds a Set backed by an in-memory tree so bootstrap can be exercised +// without the embedded FS. Each skill gets a SKILL.md plus one support file. +func mapSet(name, version string, skills ...string) Set { + files := fstest.MapFS{} + for _, s := range skills { + files["skills/"+s+"/SKILL.md"] = &fstest.MapFile{Data: []byte("---\nname: " + s + "\n---\nbody")} + files["skills/"+s+"/support.txt"] = &fstest.MapFile{Data: []byte("aux for " + s)} + } + return Set{Name: name, Version: version, Root: "skills", Skills: skills, FS: files} +} + +// TestBootstrapInstallsSkills verifies a fresh bootstrap lays every named skill +// down under skillsDir with its SKILL.md and support files intact. +func TestBootstrapInstallsSkills(t *testing.T) { + home := t.TempDir() + skills := t.TempDir() + set := mapSet("test-set", "v1", "alpha", "beta") + + bootstrap([]Set{set}, home, skills, nil) + + for _, name := range []string{"alpha", "beta"} { + md := filepath.Join(skills, name, "SKILL.md") + if _, err := os.Stat(md); err != nil { + t.Errorf("expected %s installed: %v", md, err) + } + aux := filepath.Join(skills, name, "support.txt") + if _, err := os.Stat(aux); err != nil { + t.Errorf("expected support file %s: %v", aux, err) + } + } + // State recorded so a re-run is a no-op. + if _, err := os.Stat(filepath.Join(home, stateFileName)); err != nil { + t.Errorf("expected state file written: %v", err) + } +} + +// TestBootstrapSkipsWhenAlreadyInstalled verifies a second bootstrap at the same +// version does not touch an existing (possibly user-edited) skill. +func TestBootstrapSkipsWhenAlreadyInstalled(t *testing.T) { + home := t.TempDir() + skills := t.TempDir() + set := mapSet("test-set", "v1", "alpha") + + bootstrap([]Set{set}, home, skills, nil) + + // Simulate a user edit, then re-run: the edit must survive. + md := filepath.Join(skills, "alpha", "SKILL.md") + if err := os.WriteFile(md, []byte("user edited"), 0o644); err != nil { + t.Fatalf("edit: %v", err) + } + bootstrap([]Set{set}, home, skills, nil) + + got, err := os.ReadFile(md) + if err != nil { + t.Fatalf("read: %v", err) + } + if string(got) != "user edited" { + t.Errorf("SKILL.md = %q, want user edit preserved", got) + } +} + +// TestBootstrapDoesNotClobberPreexistingSkill verifies a skill directory that +// already exists before the first bootstrap is left untouched (never overwritten). +func TestBootstrapDoesNotClobberPreexistingSkill(t *testing.T) { + home := t.TempDir() + skills := t.TempDir() + // Pre-place a user's own "alpha" skill. + if err := os.MkdirAll(filepath.Join(skills, "alpha"), 0o755); err != nil { + t.Fatal(err) + } + mine := filepath.Join(skills, "alpha", "SKILL.md") + if err := os.WriteFile(mine, []byte("mine"), 0o644); err != nil { + t.Fatal(err) + } + + bootstrap([]Set{mapSet("test-set", "v1", "alpha", "beta")}, home, skills, nil) + + got, _ := os.ReadFile(mine) + if string(got) != "mine" { + t.Errorf("preexisting alpha overwritten: %q", got) + } + // The other skill still installs. + if _, err := os.Stat(filepath.Join(skills, "beta", "SKILL.md")); err != nil { + t.Errorf("beta should install alongside preexisting alpha: %v", err) + } +} + +// TestBootstrapVersionBumpReinstallsMissing verifies bumping a set's Version +// re-triggers installation of skills that are missing, without disturbing ones +// already present. +func TestBootstrapVersionBumpReinstallsMissing(t *testing.T) { + home := t.TempDir() + skills := t.TempDir() + + bootstrap([]Set{mapSet("test-set", "v1", "alpha")}, home, skills, nil) + + // v2 adds "beta"; alpha is already present and must stay. + bootstrap([]Set{mapSet("test-set", "v2", "alpha", "beta")}, home, skills, nil) + + if _, err := os.Stat(filepath.Join(skills, "beta", "SKILL.md")); err != nil { + t.Errorf("version bump should install newly added beta: %v", err) + } +} + +// TestBootstrapEmptyHomeIsNoop verifies an unresolved home/skills dir disables +// bootstrap without panicking or erroring. +func TestBootstrapEmptyHomeIsNoop(t *testing.T) { + bootstrap([]Set{mapSet("s", "v1", "alpha")}, "", "", nil) + bootstrap([]Set{mapSet("s", "v1", "alpha")}, t.TempDir(), "", nil) +} + +// TestBootstrapEmptyVersionInstalls verifies a Set with a blank Version is not +// mistaken for "already installed" (the zero-value state lookup also equals "") +// and so its skills are installed on a fresh run. +func TestBootstrapEmptyVersionInstalls(t *testing.T) { + home := t.TempDir() + skills := t.TempDir() + + bootstrap([]Set{mapSet("test-set", "", "alpha")}, home, skills, nil) + + if _, err := os.Stat(filepath.Join(skills, "alpha", "SKILL.md")); err != nil { + t.Errorf("blank-version set should still install: %v", err) + } +} + +// TestManifestEmbedsAllSkills verifies every skill named in the real manifest is +// actually embedded (each has a SKILL.md), catching a missing-copy regression. +func TestManifestEmbedsAllSkills(t *testing.T) { + for _, set := range Manifest() { + for _, name := range set.Skills { + md := path.Join(set.Root, name, "SKILL.md") + if _, err := set.FS.Open(md); err != nil { + t.Errorf("%s/%s: embedded SKILL.md missing: %v", set.Name, name, err) + } + } + } +} + +// TestBootstrapInstallsRealManifest verifies the embedded manifest installs its +// full skill set into a temp skills dir — the end-to-end offline install path. +func TestBootstrapInstallsRealManifest(t *testing.T) { + home := t.TempDir() + skills := t.TempDir() + + Bootstrap(home, skills, nil) + + for _, set := range Manifest() { + for _, name := range set.Skills { + if _, err := os.Stat(filepath.Join(skills, name, "SKILL.md")); err != nil { + t.Errorf("skill %q not installed: %v", name, err) + } + } + } +} diff --git a/pigo/internal/builtinskills/manifest.go b/pigo/internal/builtinskills/manifest.go new file mode 100644 index 0000000..fd4fd40 --- /dev/null +++ b/pigo/internal/builtinskills/manifest.go @@ -0,0 +1,74 @@ +// Package builtinskills embeds a curated set of skills into the pigo binary and +// installs them into the user's skills directory on first run. This makes a +// baseline of workflow skills (the goal-workflow set plus a couple of standalone +// skills) available as /skill-name commands out of the box, with no manual +// `pigo install` step and no network access — the skill trees are compiled into +// the binary via //go:embed. +// +// The design is deliberately generic (a manifest of skill *sets*) so future +// collections can be added by appending a Set to Manifest and embedding their +// files, without touching the bootstrap/first-run logic in Bootstrap. +package builtinskills + +import ( + "embed" + "io/fs" +) + +// skillsFS holds the embedded skill trees under skills//. The `all:` +// prefix is required so files whose names begin with '_' or '.' (e.g. +// weather/_meta.json) are embedded too — the default pattern skips them. +// +//go:embed all:skills +var skillsFS embed.FS + +// Set is one collection of built-in skills sharing a provenance and version. +// Bootstrap installs every skill named in a Set from FS, rooted at Root (the +// path within FS holding the per-skill directories). Adding a new collection is +// a matter of appending a Set here and embedding its files — the install and +// first-run machinery is collection-agnostic. +type Set struct { + // Name identifies the collection (e.g. "goal-workflow"), used in the + // bootstrap state record and diagnostics. + Name string + // Version marks the collection's revision. It is recorded in the bootstrap + // state file; a Version newer than the recorded one re-triggers install of + // any still-missing skills on the next run. + Version string + // Root is the directory within FS that contains the per-skill subdirectories. + Root string + // Skills lists the skill directory names under Root to install. Each must be + // a "/" directory holding a SKILL.md. + Skills []string + // FS is the filesystem the skills are read from. Production sets use the + // embedded skillsFS; tests can supply an fstest.MapFS with a synthetic tree. + FS fs.FS +} + +// goalWorkflowSkills is the goal-workflow set (https://goal.rpcx.io) minus +// humanize-it (intentionally excluded), plus the two standalone skills +// architecture-diagram and weather. +var goalWorkflowSkills = []string{ + // goal-workflow core + "prd", "prd-to-spec", "to-issues", "review-it", "ship-it", + // goal-workflow bonus (humanize-it intentionally excluded) + "insight-diagram", "refactor", "modern-go", "note-it", + "code-to-spec", "smell", "loop-it", "to-design", "graph", + // standalone additions + "architecture-diagram", "weather", +} + +// Manifest is the single source of truth for the built-in skill collections +// installed on first run. It seeds one collection today; adding a Set is all +// that is needed to bundle another collection. +func Manifest() []Set { + return []Set{ + { + Name: "goal-workflow", + Version: "2026-07-24", + Root: "skills", + Skills: goalWorkflowSkills, + FS: skillsFS, + }, + } +} diff --git a/pigo/internal/builtinskills/skills/architecture-diagram/SKILL.md b/pigo/internal/builtinskills/skills/architecture-diagram/SKILL.md new file mode 100644 index 0000000..b6e5aac --- /dev/null +++ b/pigo/internal/builtinskills/skills/architecture-diagram/SKILL.md @@ -0,0 +1,163 @@ +--- +name: architecture-diagram +description: Create professional, dark-themed architecture diagrams as standalone HTML files with SVG graphics. Use when the user asks for system architecture diagrams, infrastructure diagrams, cloud architecture visualizations, security diagrams, network topology diagrams, or any technical diagram showing system components and their relationships. +license: MIT +metadata: + version: "1.0" + author: Cocoon AI (hello@cocoon-ai.com) +--- + +# Architecture Diagram Skill + +Create professional technical architecture diagrams as self-contained HTML files with inline SVG graphics and CSS styling. + +## Design System + +### Color Palette + +Use these semantic colors for component types: + +| Component Type | Fill (rgba) | Stroke | +|---------------|-------------|--------| +| Frontend | `rgba(8, 51, 68, 0.4)` | `#22d3ee` (cyan-400) | +| Backend | `rgba(6, 78, 59, 0.4)` | `#34d399` (emerald-400) | +| Database | `rgba(76, 29, 149, 0.4)` | `#a78bfa` (violet-400) | +| AWS/Cloud | `rgba(120, 53, 15, 0.3)` | `#fbbf24` (amber-400) | +| Security | `rgba(136, 19, 55, 0.4)` | `#fb7185` (rose-400) | +| Message Bus | `rgba(251, 146, 60, 0.3)` | `#fb923c` (orange-400) | +| External/Generic | `rgba(30, 41, 59, 0.5)` | `#94a3b8` (slate-400) | + +### Typography + +Use JetBrains Mono for all text (monospace, technical aesthetic): +```html + +``` + +Font sizes: 12px for component names, 9px for sublabels, 8px for annotations, 7px for tiny labels. + +### Visual Elements + +**Background:** `#020617` (slate-950) with subtle grid pattern: +```svg + + + +``` + +**Component boxes:** Rounded rectangles (`rx="6"`) with 1.5px stroke, semi-transparent fills. + +**Security groups:** Dashed stroke (`stroke-dasharray="4,4"`), transparent fill, rose color. + +**Region boundaries:** Larger dashed stroke (`stroke-dasharray="8,4"`), amber color, `rx="12"`. + +**Arrows:** Use SVG marker for arrowheads: +```svg + + + +``` + +**Arrow z-order:** Draw connection arrows early in the SVG (after the background grid) so they render behind component boxes. SVG elements are painted in document order, so arrows drawn first will appear behind shapes drawn later. + +**Masking arrows behind transparent fills:** Since component boxes use semi-transparent fills (`rgba(..., 0.4)`), arrows behind them will show through. To fully mask arrows, draw an opaque background rect (e.g., `fill="#0f172a"`) at the same position before drawing the semi-transparent styled rect on top: +```svg + + + + +``` + +**Auth/security flows:** Dashed lines in rose color (`#fb7185`). + +**Message buses / Event buses:** Small connector elements between services. Use orange color (`#fb923c` stroke, `rgba(251, 146, 60, 0.3)` fill): +```svg + +Kafka / RabbitMQ +``` + +### Spacing Rules + +**CRITICAL:** When stacking components vertically, ensure proper spacing to avoid overlaps: + +- **Standard component height:** 60px for services, 80-120px for larger components +- **Minimum vertical gap between components:** 40px +- **Inline connectors (message buses):** Place IN the gap between components, not overlapping + +**Example vertical layout:** +``` +Component A: y=70, height=60 → ends at y=130 +Gap: y=130 to y=170 → 40px gap, place bus at y=140 (20px tall) +Component B: y=170, height=60 → ends at y=230 +``` + +**Wrong:** Placing a message bus at y=160 when Component B starts at y=170 (causes overlap) +**Right:** Placing a message bus at y=140, centered in the 40px gap (y=130 to y=170) + +### Legend Placement + +**CRITICAL:** Place legends OUTSIDE all boundary boxes (region boundaries, cluster boundaries, security groups). + +- Calculate where all boundaries end (y position + height) +- Place legend at least 20px below the lowest boundary +- Expand SVG viewBox height if needed to accommodate + +**Example:** +``` +Kubernetes Cluster: y=30, height=460 → ends at y=490 +Legend should start at: y=510 or below +SVG viewBox height: at least 560 to fit legend +``` + +**Wrong:** Legend at y=470 inside a cluster boundary that ends at y=490 +**Right:** Legend at y=510, below the cluster boundary, with viewBox height extended + +### Layout Structure + +1. **Header** - Title with pulsing dot indicator, subtitle +2. **Main SVG diagram** - Contained in rounded border card +3. **Summary cards** - Grid of 3 cards below diagram with key details +4. **Footer** - Minimal metadata line + +### Component Box Pattern + +```svg + +LABEL +sublabel +``` + +### Info Card Pattern + +```html +
+
+
+

Title

+
+
    +
  • • Item one
  • +
  • • Item two
  • +
+
+``` + +## Template + +Copy and customize the template at `assets/template.html`. Key customization points: + +1. Update the `` and header text +2. Modify SVG viewBox dimensions if needed (default: `1000 x 680`) +3. Add/remove/reposition component boxes +4. Draw connection arrows between components +5. Update the three summary cards +6. Update footer metadata + +## Output + +Always produce a single self-contained `.html` file with: +- Embedded CSS (no external stylesheets except Google Fonts) +- Inline SVG (no external images) +- No JavaScript required (pure CSS animations) + +The file should render correctly when opened directly in any modern browser. diff --git a/pigo/internal/builtinskills/skills/architecture-diagram/assets/template.html b/pigo/internal/builtinskills/skills/architecture-diagram/assets/template.html new file mode 100644 index 0000000..f5b32fb --- /dev/null +++ b/pigo/internal/builtinskills/skills/architecture-diagram/assets/template.html @@ -0,0 +1,319 @@ +<!DOCTYPE html> +<html lang="en"> +<head> + <meta charset="UTF-8"> + <meta name="viewport" content="width=device-width, initial-scale=1.0"> + <title>[PROJECT NAME] Architecture Diagram + + + + +
+ +
+
+
+

[PROJECT NAME] Architecture

+
+

[Subtitle description]

+
+ + +
+ + + + + + + + + + + + + + + + + + + Users + Browser/Mobile + + + + Auth Provider + OAuth 2.0 + + + + AWS Region: us-west-2 + + + + CloudFront + CDN + + + + S3 Buckets + • bucket-one + • bucket-two + • bucket-three + OAI Protected + + + + sg-name :port + + + + Load Balancer + HTTPS :443 + + + + API Server + FastAPI :8000 + + + + Database + PostgreSQL + + + + Frontend + React + TypeScript + Additional detail + More info + domain.example.com + + + + + + HTTPS + + + + + + + OAI + + + + + TLS + + + + JWT + PKCE + + + Legend + + + Frontend + + + Backend + + + Cloud Service + + + Database + + + Security + + + Auth Flow + + + Security Group + +
+ + +
+
+
+
+

Card Title 1

+
+
    +
  • • Item one
  • +
  • • Item two
  • +
  • • Item three
  • +
  • • Item four
  • +
+
+ +
+
+
+

Card Title 2

+
+
    +
  • • Item one
  • +
  • • Item two
  • +
  • • Item three
  • +
  • • Item four
  • +
+
+ +
+
+
+

Card Title 3

+
+
    +
  • • Item one
  • +
  • • Item two
  • +
  • • Item three
  • +
  • • Item four
  • +
+
+
+ + + +
+ + diff --git a/pigo/internal/builtinskills/skills/code-to-spec/SKILL.md b/pigo/internal/builtinskills/skills/code-to-spec/SKILL.md new file mode 100644 index 0000000..9736f90 --- /dev/null +++ b/pigo/internal/builtinskills/skills/code-to-spec/SKILL.md @@ -0,0 +1,341 @@ +--- +name: code-to-spec +description: "Reverse-engineer a SPEC document from an existing project. Analyzes code, config, tests, and structure to produce a comprehensive specification. Triggers on: code-to-spec, reverse spec, generate spec, 逆向规格, 生成规格文档, 生成设计文档, 生成设计方案, extract spec, document this project, what does this project do." +user-invocable: true +--- + +# to-spec — Reverse-Engineer Project Specification + +Analyze an existing codebase and produce a structured SPEC document that captures what the project does, how it's built, and what contracts it exposes. The output is a living specification that could be used to rebuild the project from scratch or onboard new contributors. + +--- + +## When to Use + +- You want a comprehensive understanding of an existing project +- Onboarding new team members who need a high-level overview +- Documenting a project that was built without a spec +- Comparing actual implementation against intended design +- Preparing for a rewrite or major refactor +- Auditing what a project actually does vs. what people think it does + +--- + +## The Job + +1. **Scope confirmation** — ask user what to analyze (entire repo, specific directory, or specific aspect) +2. **Deep scan** — systematically read project structure, entry points, config, tests, and core logic +3. **Synthesize** — produce a structured SPEC document +4. **Review** — present to user for feedback and iteration +5. **Save** — write final SPEC to agreed location + +--- + +## Step 1: Scope Confirmation + +Before scanning, ask the user: + +``` +What should I analyze? + +A. Entire repository (recommended for small-medium projects) +B. Specific directory or module: [path] +C. Specific aspect only (e.g., API surface, data model, auth flow) + +Depth level: +1. Overview — high-level architecture + tech stack + key features (fast, ~5 min) +2. Standard — includes API contracts, data models, config, dependencies (default) +3. Deep — adds internal module interactions, error handling patterns, test coverage analysis +``` + +If the project is large (>500 files), recommend starting with Overview or a specific module. + +--- + +## Step 2: Deep Scan + +Systematically analyze the following (adapt to what exists): + +### 2.1 Project Identity +- `package.json`, `go.mod`, `Cargo.toml`, `pyproject.toml`, `pom.xml`, etc. +- README, LICENSE +- Git history (first commit date, recent activity, contributor count) + +### 2.2 Architecture +- Directory structure and organization pattern (monorepo, layered, hexagonal, etc.) +- Entry points (main files, CLI commands, server bootstrap) +- Module boundaries and dependency graph (internal) + +### 2.3 Tech Stack +- Language(s) and version constraints +- Frameworks and major libraries +- Build tools and bundlers +- Runtime requirements (Node version, Docker, etc.) + +### 2.4 Features & Behavior +- Route definitions / CLI commands / exported functions +- Business logic modules and their responsibilities +- Background jobs, cron tasks, event handlers + +### 2.5 Data Model +- Database schemas, migrations, ORMs +- Key data structures and their relationships +- State management approach + +### 2.6 API Surface +- HTTP endpoints (method, path, request/response shapes) +- GraphQL schema / gRPC protos / WebSocket events +- CLI interface (commands, flags, arguments) +- Exported library API (public functions, classes, types) + +### 2.7 Configuration & Environment +- Environment variables and their purpose +- Config files and their schema +- Feature flags, toggles + +### 2.8 External Dependencies +- Third-party services (databases, queues, APIs) +- Infrastructure requirements (cloud services, storage) +- Authentication/authorization providers + +### 2.9 Testing & Quality +- Test framework and approach (unit, integration, e2e) +- Coverage patterns (what's tested, what's not) +- Linting, formatting, type checking setup + +### 2.10 Deployment & Operations +- CI/CD configuration +- Deployment targets and strategies +- Monitoring, logging, health checks + +--- + +## Step 3: SPEC Document Structure + +Generate the SPEC with these sections. Omit sections that don't apply. + +```markdown +# SPEC: [Project Name] + +> Reverse-engineered specification — generated [date] from commit [short-hash] + +## 1. Overview + +### 1.1 Purpose +[One paragraph: what problem this project solves and for whom] + +### 1.2 Key Capabilities +- [Bullet list of what the system can do, from a user's perspective] + +### 1.3 Architecture Style +[e.g., "Monolithic Express.js API with React SPA frontend", "CLI tool with plugin system", "Microservices communicating over gRPC"] + +--- + +## 2. Tech Stack + +| Layer | Technology | Version | +|-------|-----------|---------| +| Language | ... | ... | +| Framework | ... | ... | +| Database | ... | ... | +| Build | ... | ... | +| Test | ... | ... | +| Deploy | ... | ... | + +--- + +## 3. Project Structure + +[Directory tree with annotations explaining each top-level directory's purpose] + +--- + +## 4. Data Model + +### 4.1 Core Entities +[For each entity: name, fields, relationships, constraints] + +### 4.2 State Transitions +[If applicable: lifecycle states and valid transitions] + +--- + +## 5. API Surface + +### 5.1 [Interface Type: REST / CLI / Library / etc.] + +[For each endpoint/command/function:] +| Method | Path/Command | Description | Auth | +|--------|-------------|-------------|------| +| ... | ... | ... | ... | + +### 5.2 Request/Response Schemas +[Key request/response shapes with field types] + +--- + +## 6. Configuration + +| Variable / Key | Required | Default | Description | +|---------------|----------|---------|-------------| +| ... | ... | ... | ... | + +--- + +## 7. External Dependencies + +| Service | Purpose | Failure Impact | +|---------|---------|----------------| +| ... | ... | ... | + +--- + +## 8. Business Rules & Constraints + +- [Numbered list of invariants, validation rules, and business logic constraints discovered in the code] + +--- + +## 9. Non-Functional Characteristics + +### 9.1 Performance +[Observed patterns: caching, pagination, batch processing, etc.] + +### 9.2 Security +[Auth mechanism, input validation patterns, secrets management] + +### 9.3 Error Handling +[Error strategy: custom error types, error codes, retry policies] + +--- + +## 10. Testing Strategy + +| Type | Framework | Coverage Pattern | +|------|-----------|-----------------| +| Unit | ... | ... | +| Integration | ... | ... | +| E2E | ... | ... | + +--- + +## 11. Known Gaps & Assumptions + +- [Things that are unclear from the code alone] +- [Assumptions made during analysis] +- [Areas with no tests or documentation] + +--- + +## 12. Appendix + +### A. Dependency Graph +[Key module dependencies, import relationships] + +### B. Environment Setup +[Steps to run the project locally, derived from config and scripts] +``` + +--- + +## Step 4: Review & Iteration + +After generating the SPEC, present it and ask: + +``` +SPEC generated. Please review: + +- Are there sections that need more detail? +- Are there inaccuracies I should correct? +- Should I add/remove any sections? +- Is the depth level appropriate? + +Reply OK to save, or provide feedback for iteration. +``` + +Apply feedback and re-present until user confirms. + +--- + +## Step 5: Save + +Ask user for save location: + +``` +Where should I save the SPEC? + +A. docs/SPEC.md (recommended) +B. SPEC.md (project root) +C. Custom path: [specify] +``` + +--- + +## Analysis Heuristics + +### Identifying Purpose +- Look at README first line, package description field, CLI help text +- Check the main entry point — what does it bootstrap? +- Look at test descriptions — they often describe expected behavior in plain language + +### Discovering Architecture +- Map `import`/`require` statements to build dependency graph +- Identify layers by directory naming: `controllers`, `services`, `models`, `routes`, `handlers`, `domain`, `infra` +- Check for dependency injection patterns, middleware chains, plugin registrations + +### Extracting Business Rules +- Look for validation functions, guard clauses, assertion statements +- Check error messages — they often describe what went wrong in business terms +- Examine test assertions — they encode expected behavior + +### Finding API Contracts +- Route registrations (Express: `app.get()`, FastAPI: `@app.get()`, Go: `mux.HandleFunc()`) +- OpenAPI/Swagger files if present +- Request validation schemas (Joi, Zod, Pydantic, struct tags) +- CLI flag/argument definitions (cobra, argparse, yargs) + +### Detecting Data Models +- ORM model definitions (Prisma, SQLAlchemy, GORM, TypeORM) +- Migration files (in chronological order) +- Type/interface definitions for core domain objects +- Database seed files + +--- + +## Edge Cases + +| Scenario | Handling | +|----------|----------| +| Project has no README or documentation | Note this in "Known Gaps"; infer purpose from code | +| Monorepo with multiple services | Ask user which service(s) to analyze; produce one SPEC per service or a unified SPEC with clear boundaries | +| Project uses code generation | Document the generated code's purpose but focus on the source of truth (schemas, proto files, templates) | +| Legacy project with mixed patterns | Document all observed patterns, note inconsistencies in "Known Gaps" | +| Project is a library (no runtime) | Focus on exported API surface, type contracts, and usage patterns from tests | +| Incomplete or broken code | Document what exists, mark broken/incomplete areas explicitly | +| Project >1000 files | Start with entry points and trace key flows; don't exhaustively read every file | +| Multiple languages in one repo | Document each language's role and how they interact | + +--- + +## Quality Criteria + +A good reverse-engineered SPEC should pass these checks: + +- [ ] A developer unfamiliar with the project could understand its purpose in 60 seconds +- [ ] The tech stack section is complete enough to set up a dev environment +- [ ] API contracts are specific enough to write a client against +- [ ] Data models are complete enough to recreate the schema +- [ ] Business rules are explicit (not buried in "see code") +- [ ] Known gaps are honestly listed (don't invent what you can't determine) +- [ ] The SPEC matches the actual code (not aspirational documentation) + +--- + +## Anti-Patterns to Avoid + +- **Don't invent intent.** If you can't determine WHY something exists, say so. Don't fabricate rationale. +- **Don't copy code into the SPEC.** Describe behavior and contracts, don't paste implementations. +- **Don't include transient state.** The SPEC describes the system's design, not its current runtime state. +- **Don't over-specify internals.** Focus on boundaries, contracts, and behavior. Internal implementation details belong in code comments, not specs. +- **Don't assume the README is accurate.** READMEs often lag behind code. Verify claims against actual implementation. diff --git a/pigo/internal/builtinskills/skills/graph/SKILL.md b/pigo/internal/builtinskills/skills/graph/SKILL.md new file mode 100644 index 0000000..e8c1e89 --- /dev/null +++ b/pigo/internal/builtinskills/skills/graph/SKILL.md @@ -0,0 +1,329 @@ +--- +name: graph +description: "Graph engineering for parallel task execution: convert a task, PRD, SPEC, or issue set into a dependency graph (DAG), layer it into supersteps, then implement each independent node concurrently with subagents — each node runs /goal → /review-it → /ship-it in an isolated git worktree, with a fan-in barrier between waves. Triggers on: graph, graph engineering, build a graph, task graph, dependency graph, DAG, parallel implement, 并发实现, 并行实现, 任务图, 把任务变成图, fan-out fan-in, superstep, dynamic workflow." +user-invocable: true +allowed-tools: + - Bash(git:*) + - Bash(gh:*) + - Bash(cat:*) + - Bash(mkdir:*) + - Bash(grep:*) + - Bash(python3:*) +--- + +# graph — Task/PRD to Parallel Execution Graph + +Turn a task (or PRD / SPEC / issue set) into a **directed acyclic graph** of work units, layer it into **supersteps (waves)**, and implement each wave's independent nodes **concurrently** using subagents. Each node runs the full `/goal → /review-it → /ship-it` pipeline inside its **own git worktree**, so parallel nodes never clobber each other's working tree. Between waves, a **fan-in barrier** merges results and re-plans the next wave. + +This is the parallel sibling of `/loop-it`. `/loop-it` is strictly sequential (one worktree, one issue at a time). `/graph` fans out every independent node in a wave at once. + +--- + +## Mental Model (borrowed from LangGraph / graph engineering) + +| Concept | Here | +|---------|------| +| **Node** | One implementable unit of work (an issue / subtask) | +| **Edge** | A dependency: `B depends on A` → edge `A → B` | +| **Superstep / wave** | A set of nodes whose deps are all satisfied — run concurrently | +| **Fan-out** | Dispatch one subagent per node in the current wave | +| **Fan-in (barrier)** | Wait for **all** nodes in the wave before starting the next | +| **State channel** | `.graph_state` — shared checkpoint, rewritten between waves (resume source) | +| **Live tracker** | `graph.html` — Claude-style light-theme dashboard, re-rendered from `.graph_state` at every checkpoint | +| **Dynamic re-plan** | After a wave, revise the graph if new work/deps emerged | + +**Core principle:** Independent nodes in the same wave have *no shared state and no ordering dependency*, so they can run in true parallel. Dependencies define the *only* ordering. Everything else runs at once. + +--- + +## Overview + +``` +Input (task / PRD / SPEC / issues) + │ + ▼ +1. Decompose into nodes ─────────► nodes = {id, title, deps, criteria, scope} + │ + ▼ +2. Build DAG + validate ─────────► detect cycles, orphan deps + │ + ▼ +3. Topological layering ─────────► waves = [[n1,n2,n3], [n4,n5], [n6]] + │ + ▼ +4. Render graph + confirm with user + │ + ▼ (write .graph_state + graph.html — open graph.html to watch live) +┌──────────── per wave (superstep) ────────────┐ +│ │ +│ FAN-OUT: 1 subagent per node (parallel) │ +│ each subagent, in its own git worktree: │ +│ /goal (inline implement) → /review-it │ +│ → /ship-it │ +│ │ +│ FAN-IN barrier: wait for ALL nodes │ +│ integrate, update .graph_state │ +│ re-render graph.html │ +│ re-plan next wave if graph changed │ +│ │ +└───────────────────────────────────────────────┘ + │ + ▼ +All waves done → final summary +``` + +--- + +## Step 1: Locate & Decompose Input + +Accept any of: a free-form task description, a PRD/SPEC file, or an existing issue set (GitHub / local `.md` / iCafe). + +- **PRD/SPEC** → reuse `/to-issues` decomposition rules (one node per User Story; split large, merge tiny). +- **Existing issues** → each issue is a node; parse dependencies from issue bodies (`Depends on: #3`, `Dependencies: #3, #5`). +- **Free-form task** → break into the smallest independently-shippable units yourself. + +Each node MUST have: + +``` +Node #N + title: short imperative title + deps: [list of node ids] or [] + criteria: acceptance criteria (checklist) — how the subagent knows it's done + type: backend | frontend | fullstack | ui | infra | docs + scope_hint: which files/dirs this node is expected to touch (for conflict analysis) +``` + +`scope_hint` matters: two nodes with no dependency edge but overlapping file scope are **not** truly independent — see Step 3. + +--- + +## Step 2: Build the DAG & Validate + +Construct edges from `deps`. Then validate: + +| Check | Action on failure | +|-------|-------------------| +| **Cycle** (`A → B → A`) | Print `⚠️ 循环依赖: #A ↔ #B`. Break by node id order, warn user, ask to confirm or fix. | +| **Dangling dep** (`#7 depends on #99`, no such node) | Print warning, drop the phantom edge. | +| **Scope collision** (two dep-free nodes edit same files) | Add a *soft edge* to serialize them (lower id first), OR flag for user. Never let two parallel worktrees fight over the same files. | + +**Hot-file exception:** A shared *wiring* file that nearly every node must touch (e.g. `router.go`, `main.go`, `mod.rs`, a DI container, an `__init__` re-export) does NOT count as a scope collision — treating it as one would serialize the entire graph into a chain. For such files, assume append-only edits merge cleanly, and prefer one of: (a) designate a single node that *owns* wiring and have others expose a registration hook, or (b) do a tiny follow-up "wire everything" node in the last wave. Reserve the collision rule for nodes that edit the *same logic* in the same file (e.g. two handlers rewriting the same function). + +--- + +## Step 3: Topological Layering into Waves + +Compute waves via Kahn's algorithm: + +1. **Wave 0** = all nodes with `deps == []` and no scope collision among themselves. +2. Remove wave-0 nodes; **Wave 1** = nodes whose deps are now all satisfied. +3. Repeat until all nodes placed. +4. Within a wave, if two nodes edit the **same logic in the same file** (real collision, per the hot-file exception in Step 2), push the higher-id one to the next wave. Bare wiring-file overlap does not trigger this. + +**ID conventions (used consistently):** lower id wins — cycles break by lowest id first (Step 2), and scope collisions serialize with the lower id first (higher id deferred to the next wave). + +Print the layered plan: + +``` +📊 Graph: 6 nodes, 3 waves + +Wave 0 (parallel ×3): #1 db schema #2 config loader #3 logging util +Wave 1 (parallel ×2): #4 API handler (deps #1) #5 CLI flags (deps #2) +Wave 2 (parallel ×1): #6 integration (deps #4,#5) + +Max parallelism: 3 subagents in Wave 0. +``` + +Also emit a Mermaid diagram for the user: + +```` +```mermaid +graph LR + n1[#1 db schema] --> n4[#4 API handler] + n2[#2 config loader] --> n5[#5 CLI flags] + n3[#3 logging util] + n4 --> n6[#6 integration] + n5 --> n6 +``` +```` + +**Wait for user confirmation** before dispatching any subagent. Let them adjust nodes, deps, or the max-parallelism cap. + +--- + +## Step 4: Pre-flight Checks + +Before the first wave (same spirit as `/loop-it`): + +```bash +git rev-parse --is-inside-work-tree # in a repo? +git status --porcelain # clean tree? (dirty → stash/abort) +git branch --show-current # on main/master? +git ls-remote --heads origin # remote reachable? +gh auth status # if shipping to GitHub +``` + +Any hard failure → print the error and stop. Confirm a **max concurrency cap** with the user (default 3–4 parallel subagents; more risks rate limits and review noise). + +Then **initialize the state channel + live tracker** (do this once, right after the plan is confirmed and before the first fan-out): + +```bash +# 1. Write the initial checkpoint (all nodes pending, current_wave 0). +cat > .graph_state <<'JSON' +{ "version": 1, "task": "...", "repo": "owner/repo", + "waves": [[1,2,3],[4,5],[6]], "current_wave": 0, + "nodes": { "1": {"title":"...","deps":[],"status":"pending","wave":0}, ... } } +JSON + +# 2. Keep it out of git. +grep -qxF '.graph_state' .gitignore || printf '.graph_state\ngraph.html\n' >> .gitignore + +# 3. Render the Claude-style light-theme dashboard. +python3 skills/graph/scripts/render_graph_html.py .graph_state graph.html +``` + +Tell the user: **open `graph.html` in a browser** — it auto-refreshes every 5s, so it tracks execution live (waves, node statuses, progress bar, and a Mermaid DAG colored by status). Re-run the render command at every checkpoint (see Step 5b) to push updates. + +--- + +## Step 5: Execute Wave by Wave (fan-out → fan-in) + +For each wave, in order: + +### 5a. FAN-OUT — one subagent per node, in parallel + +**Dispatch all nodes of the wave in a single response** (multiple Agent/subagent calls in one message = concurrent). Each subagent works in its **own git worktree** so parallel file edits never collide: + +```bash +# The orchestrator creates a worktree per node BEFORE dispatching: +git worktree add -b feat/node-{N}-{slug} ../.graph-worktrees/node-{N} main +``` + +Each subagent receives a **self-contained** prompt (it does NOT inherit orchestrator context): + +```markdown +You are implementing ONE node of a task graph, working in an ISOLATED git worktree. + +Worktree: ../.graph-worktrees/node-{N} (already created on branch feat/node-{N}-{slug}) +Node #{N}: {title} +Type: {type} +Scope: {scope_hint} — stay within these files; do not touch other nodes' scope + +Acceptance criteria (all must pass): +- [ ] {criterion 1} +- [ ] {criterion 2} + +Context (deps already merged into main, pull first): +{summaries of dependency nodes' outputs, or the referenced PRD/SPEC excerpt} + +Your pipeline (run all three, in order): +1. IMPLEMENT (inline /goal): read the node + any referenced PRD/SPEC, read adjacent + code, implement to satisfy EVERY acceptance criterion, run build + tests + lint + (e.g. go build ./... && go vet ./... && go test ./...). Iterate until all green. +2. REVIEW (/review-it): run code review on your changes, apply accepted findings, + re-run focused tests, repeat until review is clean (max 2 rounds). +3. SHIP (/ship-it): commit (message references the node/issue), push branch, + create PR, merge, close the issue. + +Constraints: +- Work ONLY inside your worktree. Do NOT edit files outside {scope_hint}. +- Do NOT try to call `goal` via the Skill tool (it's a UI command, not a skill) — + "implement" means you write the code yourself. /review-it and /ship-it ARE skills. +- If you cannot satisfy a criterion, STOP and report what's blocking — don't fake it. + +Return: node id, PASS/FAIL, PR/commit refs, files changed, and — if you discovered new required work or a dependency the graph didn't capture — a `NEW_WORK:` line describing it (title + which nodes it blocks). Emit `NEW_WORK: none` if there's nothing. +``` + +> **Why worktrees, not branches alone:** `/goal` mutates the working tree. Two subagents editing the same checkout would corrupt each other. A worktree per node gives each its own filesystem checkout on its own branch — that's what makes the wave genuinely parallel and safe. + +### 5b. FAN-IN — barrier, integrate, re-plan + +Wait for **every** subagent in the wave to return (BSP barrier — the next wave cannot start until this one commits). Then: + +1. Read each subagent's summary. Mark node `shipped` or `failed`. +2. `git checkout main && git pull` — dependency outputs are now on main for the next wave. +3. Remove finished worktrees: `git worktree remove ../.graph-worktrees/node-{N}` (keep failed ones for investigation). +4. Write checkpoint to `.graph_state`, then re-render the tracker: + `python3 skills/graph/scripts/render_graph_html.py .graph_state graph.html` (the open `graph.html` picks it up on its next auto-refresh). +5. **Dynamic re-plan** (LangGraph-style conditional edge): scan each subagent's `NEW_WORK:` line. If any is not `none`, add the new node(s)/edge(s) and re-layer the *remaining* nodes before starting the next wave. Show the user the delta. +6. If any node in the wave **failed**, mark all nodes that depend on it as `blocked` and skip them (their inputs aren't ready). + +Proceed to the next wave. + +--- + +## State File: `.graph_state` (+ live tracker `graph.html`) + +`.graph_state` lives at the repo root and **must be in `.gitignore`**. It's the single source of truth: checkpoint it after every wave so a crash resumes at the wave boundary, and re-render `graph.html` from it so the browser dashboard stays live. `graph.html` is a *derived* view — never hand-edit it; regenerate it from `.graph_state`. + +```json +{ + "version": 1, + "updated_at": "2026-07-21T10:30:00Z", + "task": "Add user auth", + "repo": "owner/repo", + "waves": [[1, 2, 3], [4, 5], [6]], + "current_wave": 1, + "nodes": { + "1": { "title": "db schema", "deps": [], "status": "shipped", "branch": "feat/node-1-db-schema", "pr": 43, "wave": 0 }, + "2": { "title": "config loader", "deps": [], "status": "shipped", "wave": 0 }, + "3": { "title": "logging util", "deps": [], "status": "failed", "wave": 0, "error": "test TestLog failed", "attempts": 2 }, + "4": { "title": "API handler", "deps": [1], "status": "in_progress", "wave": 1 }, + "6": { "title": "integration", "deps": [4, 5], "status": "blocked", "wave": 2, "reason": "depends on #3 (failed)" } + } +} +``` + +Status values: `pending | in_progress | shipped | failed | blocked | skipped`. Each node carries `title` + `deps` so `graph.html` can draw the DAG and cards straight from the checkpoint. + +Render the tracker any time with: + +```bash +python3 skills/graph/scripts/render_graph_html.py .graph_state graph.html +``` + +On resume: read `.graph_state`, skip `shipped`, ask about `failed` (retry/skip), re-derive remaining waves, and re-render `graph.html`. + +--- + +## Safety Guards + +- **Worktree isolation is mandatory** — never run two parallel `/goal` sessions in the same checkout. +- **Fan-in barrier is mandatory** — never start wave N+1 before every node in wave N returns and merges. +- **Scope collisions serialize** — dep-free nodes touching the same files go in different waves. +- **Never skip /review-it** before `/ship-it`. +- **Cap concurrency** — default 3–4; more invites rate limits and merge contention. +- **Never force-push to main.** Each node ships via its own branch/PR. +- **Failed node blocks its dependents** — don't ship on top of unmet inputs. +- **Max retries per node** — reuse `/loop-it`'s error classes; don't loop forever. +- **Confirm the plan** before the first fan-out. + +--- + +## Common Mistakes + +| Mistake | Fix | +|---------|-----| +| Dispatching subagents in separate responses | One response, multiple calls = parallel. Separate = sequential. | +| No worktree → parallel edits corrupt the tree | One `git worktree` per node. | +| Two "independent" nodes edit the same file | Add a soft edge; put them in different waves. | +| Starting the next wave before all nodes merge | Enforce the fan-in barrier. | +| Over-decomposing into 20 trivial nodes | Merge tiny units; a node should be a meaningful shippable unit. | +| Ignoring a failed node's dependents | Mark them `blocked`, skip them. | + +--- + +## Relationship to Other Skills + +``` +/prd → /prd-to-spec → /to-issues ─┬─► /loop-it (sequential: one node at a time) + └─► /graph (parallel: whole wave at once) + │ + each node: inline /goal → /review-it → /ship-it (in its own worktree) +``` + +- **`/to-issues`** — decomposition rules reused for building nodes. +- **`/loop-it`** — sequential counterpart; use it when nodes heavily share files or serial safety matters. +- **`/graph`** — this skill; use it when the DAG has genuine parallelism (independent subsystems). +- **`/review-it`, `/ship-it`** — real skills each node's subagent invokes. +``` diff --git a/pigo/internal/builtinskills/skills/graph/scripts/render_graph_html.py b/pigo/internal/builtinskills/skills/graph/scripts/render_graph_html.py new file mode 100644 index 0000000..3db25a8 --- /dev/null +++ b/pigo/internal/builtinskills/skills/graph/scripts/render_graph_html.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +"""Render a Claude-style light-theme graph.html dashboard from a .graph_state state file. + +Usage: + render_graph_html.py [.graph_state] [graph.html] + +Defaults to reading ./.graph_state and writing ./graph.html. +Called by the /graph skill at every checkpoint (initial plan + each fan-in barrier), +so opening graph.html in a browser (it self-refreshes) tracks execution live. +No third-party dependencies — stdlib only. +""" +import html +import json +import sys +from datetime import datetime + +STATUS = { + "pending": ("Pending", "#8C8579", "#EFECE3"), + "in_progress": ("In Progress", "#CC785C", "#F7E9E2"), + "shipped": ("Shipped", "#3D7A5A", "#DFEEE4"), + "failed": ("Failed", "#B54A3E", "#F6DEDA"), + "blocked": ("Blocked", "#9A6C3A", "#F2E6D4"), + "skipped": ("Skipped", "#8C8579", "#EFECE3"), +} + + +def esc(s): + return html.escape(str(s if s is not None else "")) + + +def node_card(nid, n): + st = n.get("status", "pending") + label, fg, bg = STATUS.get(st, STATUS["pending"]) + deps = n.get("deps") or [] + deps_str = ", ".join(f"#{d}" for d in deps) if deps else "no deps" + meta = [] + if n.get("pr"): + meta.append(f'PR #{esc(n["pr"])}') + if n.get("branch"): + meta.append(f'{esc(n["branch"])}') + if n.get("attempts"): + meta.append(f'attempt {esc(n["attempts"])}') + meta_html = " · ".join(meta) + err = f'
{esc(n["error"])}
' if n.get("error") else "" + return f""" +
+
+ #{esc(nid)} + {label} +
+
{esc(n.get('title','(untitled)'))}
+
{esc(deps_str)}
+ {f'
{meta_html}
' if meta_html else ''} + {err} +
""" + + +def mermaid(state): + lines = ["graph LR"] + nodes = state.get("nodes", {}) + for nid, n in nodes.items(): + t = n.get("title", "") + lines.append(f' n{nid}["#{nid} {t}"]') + for nid, n in nodes.items(): + for d in (n.get("deps") or []): + lines.append(f" n{d} --> n{nid}") + # color by status + for st, (_, fg, bg) in STATUS.items(): + ids = [f"n{nid}" for nid, n in nodes.items() if n.get("status") == st] + if ids: + lines.append(f" classDef {st} fill:{bg},stroke:{fg},color:#33312B;") + lines.append(f" class {','.join(ids)} {st};") + return "\n".join(lines) + + +def render(state): + nodes = state.get("nodes", {}) + total = len(nodes) + counts = {k: 0 for k in STATUS} + for n in nodes.values(): + counts[n.get("status", "pending")] = counts.get(n.get("status", "pending"), 0) + 1 + shipped = counts.get("shipped", 0) + pct = int(shipped / total * 100) if total else 0 + waves = state.get("waves", []) + cur = state.get("current_wave", 0) + + legend = "".join( + f'{label}' + for label, fg, bg in STATUS.values() + ) + + wave_html = "" + for wi, wave in enumerate(waves): + state_cls = "cur" if wi == cur else ("done" if wi < cur else "future") + cards = "".join(node_card(str(nid), nodes.get(str(nid), {"title": f"#{nid}"})) for nid in wave) + wave_html += f""" +
+

Wave {wi} ×{len(wave)} parallel + {'running' if wi == cur else ''}

+
{cards}
+
""" + + stat = lambda k: f'{counts.get(k,0)} {STATUS[k][0].lower()}' + stats = " · ".join(stat(k) for k in ["shipped", "in_progress", "failed", "blocked", "skipped", "pending"]) + + return f""" + + + + + +graph · {esc(state.get('task','execution'))} + + + + +
+
+

{esc(state.get('task','Task Graph Execution'))}

+
{esc(state.get('repo',''))} · wave {cur} of {max(len(waves)-1,0)} · updated {esc(state.get('updated_at',''))}
+
+
{shipped}/{total} shipped ({pct}%)  —  {stats}
+
{legend}
+
+
{esc(mermaid(state))}
+ {wave_html} +
Auto-refreshes every 5s · generated by /graph from .graph_state
+
+ + +""" + + +def main(): + src = sys.argv[1] if len(sys.argv) > 1 else ".graph_state" + dst = sys.argv[2] if len(sys.argv) > 2 else "graph.html" + with open(src, encoding="utf-8") as f: + state = json.load(f) + state.setdefault("updated_at", datetime.now().isoformat(timespec="seconds")) + with open(dst, "w", encoding="utf-8") as f: + f.write(render(state)) + print(f"wrote {dst} from {src}") + + +if __name__ == "__main__": + main() diff --git a/pigo/internal/builtinskills/skills/insight-diagram/SKILL.md b/pigo/internal/builtinskills/skills/insight-diagram/SKILL.md new file mode 100644 index 0000000..4540be7 --- /dev/null +++ b/pigo/internal/builtinskills/skills/insight-diagram/SKILL.md @@ -0,0 +1,273 @@ +--- +name: insight-diagram +description: 为任意项目生成 UML 图、架构图和流程图。分析代码库后让用户选择要生成的图表类型,使用 architecture-diagram skill 渲染为 HTML+SVG,保存到 docs/ 目录。适用于任何软件项目的文档可视化。 +--- + +# Insight Diagram — 项目图表生成技能 + +分析任意代码库,自动生成 UML 14种图 + 架构图 + 流程图,使用 `/architecture-diagram` 渲染为 HTML+SVG。 + +## 图表分类与清单 + +### 结构性图形 (Structural Diagrams — 静态) +描述系统的物理组成和静态结构。 + +| 编号 | 图表类型 | 英文标识 | 关注点 | +|------|---------|---------|--------| +| 1 | 系统架构图 | architecture | 组件关系、全局视角(非UML,最常用) | +| 2 | 类图 | class | 定义类、属性、操作及关系 | +| 3 | 对象图 | object | 特定时刻的对象实例及其关系 | +| 4 | 组件图 | component | 系统组件及其依赖关系 | +| 5 | 部署图 | deployment | 物理硬件、节点及软件部署 | +| 6 | 包图 | package | 将模型元素分组组织 | +| 7 | 复合结构图 | composite-structure | 类的内部结构 | +| 8 | 剖面图 | profile | 扩展UML元模型、自定义构造型 | + +### 行为性图形 (Behavioral Diagrams — 动态) +描述系统与外部参与者或系统内部的交互过程。 + +| 编号 | 图表类型 | 英文标识 | 关注点 | +|------|---------|---------|--------| +| 9 | 流程图 | flowchart | 主流程与分支(非UML,最常用) | +| 10 | 用例图 | usecase | 从用户角度展示系统功能 | +| 11 | 活动图 | activity | 过程的流程或步骤 | +| 12 | 状态机图 | state-machine | 对象生命周期的状态变迁 | +| 13 | 序列图 | sequence | 按时间顺序展示对象间交互 | +| 14 | 通信图 | communication | 侧重于对象间的组织关系 | +| 15 | 定时图 | timing | 侧重于状态变化的时间约束 | +| 16 | 交互概览图 | interaction-overview | 结合活动图和时序图 | +| 17 | 泳道图 | swimlane | 跨组件/角色职责流程(活动图变体) | + +## 示例参考 + +本技能的 `examples/` 目录包含 13 个已完成的图表 HTML 文件,作为视觉样式和内容结构的参考模板。**生成任何图表前,必须先阅读对应的示例文件**,以确保风格一致、结构规范。 + +### 示例文件清单 + +| 文件 | 图表类型 | 英文标识 | +|------|---------|---------| +| `examples/architecture.html` | 系统架构图 | architecture | +| `examples/class.html` | 类图 | class | +| `examples/object.html` | 对象图 | object | +| `examples/component.html` | 组件图 | component | +| `examples/deployment.html` | 部署图 | deployment | +| `examples/flowchart.html` | 流程图 | flowchart | +| `examples/usecase.html` | 用例图 | usecase | +| `examples/activity.html` | 活动图 | activity | +| `examples/sequence.html` | 序列图 | sequence | +| `examples/communication.html` | 通信图 | communication | +| `examples/dfd.html` | 数据流图 | dfd | +| `examples/interaction-overview.html` | 交互概览图 | interaction-overview | +| `examples/swimlane.html` | 泳道图 | swimlane | + +### 参考规则 + +1. **生成前必读**: 调用 `/architecture-diagram` 前,先用 Read 工具阅读对应类型的示例文件,从中提取: + - SVG 布局策略(节点间距、分组方式、箭头走向) + - 节点样式层级(核心节点 accent 高亮、普通节点实线边框、可选节点虚线边框) + - 标注风格(阶段标签、Legend 图例、卡片摘要) + - 信息密度(每个节点显示多少字段/属性) + +2. **结构对齐**: 生成的图表应与示例保持相同的结构层次: + - 页面顶部:标题 + 副标题 + 图表类型说明 + - 中间主体:SVG 图表区域(带浅色边框容器) + - 底部:信息摘要卡片 + 页脚 + +3. **内容替换而非照搬**: 示例中的业务数据(NovaShield 风控系统)是虚构的参考案例,生成时需替换为目标项目的真实架构信息。只参考布局和样式,不复制业务内容。 + +4. **无对应示例的类型**: 对于包图 (package)、复合结构图 (composite-structure)、剖面图 (profile)、状态机图 (state-machine)、定时图 (timing) 这 5 种没有示例文件的图表类型,参考最相近的已有示例(如包图参考组件图,状态机图参考活动图),并沿用相同的视觉语言。 + +## 执行流程 + +### 步骤 1:分析代码库 + +读取项目关键文件,提取架构信息: + +1. 读取项目根目录的 `CLAUDE.md`(如存在)获取项目概览 +2. 读取各子目录的 `CLAUDE.md`(如存在)获取模块细节 +3. 用 Glob 扫描源码文件结构(`**/*.go`, `**/*.py`, `**/*.ts` 等) +4. 读取入口文件(`main.go`, `app.py`, `index.ts` 等)识别顶层组件 +5. 用 Grep 搜索关键模式:接口定义、函数签名、依赖注入、配置项 + +从以上信息中提炼出: +- **组件清单**: 服务、模块、外部依赖 +- **关系图**: 谁调用谁、谁依赖谁、数据流向 +- **核心类型**: 结构体/类、接口、枚举 +- **流程**: 主业务流程、异常处理流程 +- **部署**: 进程、中间件、外部服务 + +### 步骤 2:选择图表 + +使用 AskUserQuestion 让用户选择要生成的图表(multiSelect: true),分4组展示: + +**第1组 — 结构性图形(静态):** +- 系统架构图 (architecture) +- 类图 (class) +- 对象图 (object) +- 组件图 (component) + +**第2组 — 结构性图形续 + 部署:** +- 部署图 (deployment) +- 包图 (package) +- 复合结构图 (composite-structure) +- 剖面图 (profile) + +**第3组 — 行为性图形(动态):** +- 流程图 (flowchart) +- 用例图 (usecase) +- 活动图 (activity) +- 状态机图 (state-machine) + +**第4组 — 交互图 + 常用非UML:** +- 序列图 (sequence) +- 通信图 (communication) +- 交互概览图 (interaction-overview) +- 泳道图 (swimlane) +- 全部生成 (all) + +默认推荐:architecture + sequence + flowchart + +### 步骤 3:逐个生成 + +对每个选中的图表类型: + +1. **先读示例**: 用 Read 工具阅读 `examples/<标识>.html`(如 `examples/architecture.html`),提取布局模式、节点样式、标注方式 +2. 根据步骤 1 提取的架构信息,整理出该图表应展示的元素和关系 +3. 调用 `/architecture-diagram` skill,传入图表类型、标题、内容描述、输出路径,**必须指定 light 风格** +4. 输出文件保存到 `docs/<标识>.html`(如 `docs/architecture.html`) +5. **生成后必须 review**: 运行几何校验脚本,按结果修正后再继续下一张(见下方「步骤 3.5」) +6. 简要报告完成状态 + +**生成规则:** +- **风格**: 必须使用 light Claude 风格(暖白背景 #FAF9F6、terracotta/sage/plum/rose 配色、Inter 字体、白色卡片容器),与 Anthropic Claude 品牌视觉一致 +- **防遮盖**: 所有 SVG 元素(节点、箭头、标签)不得互相遮盖。具体做法: + - 计算每个元素的边界框,确保无重叠 + - 箭头绘制在节点下方(SVG 中先画箭头再画节点) + - 节点间留足间距(垂直最少 40px,水平最少 30px) + - 文字不超出所在节点边界,超长文字截断或换行 + - 连接线的标签放置在线段中点偏移处,避免覆盖线段或节点 + - 如果元素过多导致图表拥挤,拆分为多个子图或缩小元素尺寸 + +批量生成顺序(宏观→微观): +architecture → component → deployment → package → composite-structure → profile → class → object → usecase → flowchart → activity → state-machine → swimlane → sequence → communication → timing → interaction-overview + +### 步骤 3.5:几何 review(每张图生成后必做) + +生成的 SVG 常见三类几何缺陷,必须用脚本逐张校验并修正: + +```bash +python3 skills/insight-diagram/scripts/review_svg.py docs/<标识>.html --min-gap 8 +# 批量: python3 skills/insight-diagram/scripts/review_svg.py docs/*.html --min-gap 8 +``` + +脚本检查(与三条核心要求一一对应): + +1. **箭头落点**:每个带箭头的端点必须恰好落在目标框/椭圆/菱形的**边缘线**上(容差 6px)。 + - `ERROR 深入框内`:端点穿入框内部 >8px → 缩短连线,让它止于边缘。 + - `WARNING 空接`:端点悬空、距最近框 >8px 且不汇入任何其它连线 → 把端点对齐到框边或汇合点。 + - 合法情形:端点落在框边、生命线、或与另一条连线交汇(分支/汇聚)——脚本不会误报。 +2. **框重叠**(`ERROR`):非嵌套的两个框在水平、垂直两个方向都有交叠 → 必须移开其中一个。嵌套(一个完全包住另一个,如分组边界框包子节点)是允许的。 +3. **框间距**(`WARNING`):投影相邻的两框净间距 < `--min-gap`(默认 8px)→ 拉开距离。 + +处理原则: +- **ERROR 必须修复**后再进入下一张;修完重跑脚本确认归零。 +- **WARNING 逐条核对**:序列图生命线底部的消息、泳道边界、紧贴的分组等可能是设计本意,确认无误可保留;其余应调整坐标。 +- 修正方式是直接编辑 `docs/<标识>.html` 里对应的 `//` 坐标,而非重新生成整张图。 +- 退出码:有 ERROR 返回 1,干净返回 0;CI 中可加 `--strict` 让 WARNING 也阻断。 + +### 步骤 4:报告 + +全部完成后输出: +- 生成的文件列表 +- 每个图表的简要描述 + +## 各图表的内容指南 + +### 系统架构图 (architecture) — 非UML,最常用 +- 展示系统顶层组件及其连接关系 +- 区分内部模块与外部依赖 +- 标注核心数据流方向 + +### 类图 (class) +- 核心类型为类节点(名称+字段+方法) +- 继承、组合、依赖关系 +- 接口与实现分离 +- 限制在 10-15 个核心类型 + +### 对象图 (object) +- 选取一个典型运行时场景 +- 展示对象实例及其属性值 +- 对象间的链接关系 + +### 组件图 (component) +- 每个组件为一个节点 +- 箭头表示依赖/调用方向 +- 标注接口名称 + +### 部署图 (deployment) +- 物理节点(服务器、容器、Serverless) +- 中间件(消息队列、缓存、数据库) +- 外部服务(第三方 API) +- 标注通信协议 + +### 包图 (package) +- 按模块/命名空间分组 +- 包间依赖关系 +- 体现分层架构 + +### 复合结构图 (composite-structure) +- 类/组件的内部结构 +- 部件(Part)与连接器(Connector) +- 端口(Port)与接口 + +### 剖面图 (profile) +- 自定义构造型(Stereotype) +- 扩展元模型的标签定义(Tagged Values) +- 领域特定建模约束 + +### 流程图 (flowchart) — 非UML,最常用 +- 主流程 + 关键分支 +- 失败/异常路径 +- 起止节点清晰 + +### 用例图 (usecase) +- 参与者(人/外部系统) +- 用例椭圆 +- include/extend 关系 + +### 活动图 (activity) +- 阶段/步骤为活动节点 +- 并行分支用 fork/join +- 决策点用菱形 + +### 状态机图 (state-machine) +- 对象的关键状态 +- 触发状态变迁的事件 +- 动作/守卫条件 +- 初始态和终态 + +### 序列图 (sequence) +- 参与者为纵向生命线 +- 水平箭头为消息调用 +- 标注关键返回值 +- 关注 2-5 个核心交互场景 + +### 通信图 (communication) +- 组件为节点,消息为连线 +- 标注消息序号 +- 强调协作关系而非时序 + +### 定时图 (timing) +- 时间轴横向展开 +- 状态变化的时间约束 +- 持续时间标注 + +### 交互概览图 (interaction-overview) +- 控制流节点内嵌交互片段 +- 展示条件分支和循环 +- 宏观概览各交互场景 + +### 泳道图 (swimlane) — 活动图变体 +- 按组件/角色分泳道 +- 流程步骤在对应泳道内 +- 跨泳道箭头表示交互 diff --git a/pigo/internal/builtinskills/skills/insight-diagram/references/extraction-strategy.md b/pigo/internal/builtinskills/skills/insight-diagram/references/extraction-strategy.md new file mode 100644 index 0000000..f85dab1 --- /dev/null +++ b/pigo/internal/builtinskills/skills/insight-diagram/references/extraction-strategy.md @@ -0,0 +1,62 @@ +# 图表类型内容提取策略 + +每种图表需要从代码库中提取不同维度的信息。以下是通用的提取策略,适用于任何语言/框架的项目。 + +## 通用提取规则 + +### 组件识别 +- 入口文件中的初始化/注册代码 +- 依赖注入容器(Wire, Spring, 等) +- 包/模块的公开接口 +- 配置文件中引用的外部服务 + +### 关系识别 +- import/require 语句 +- 函数调用链(谁调用了谁) +- 接口实现关系 +- 事件发布/订阅 + +### 数据流识别 +- 函数参数和返回值 +- 消息队列的 topic/producer/consumer +- API endpoint 的 request/response +- 数据库读写操作 + +### 流程识别 +- 主循环 / 事件循环 +- 中间件链 / handler 链 +- 状态机转换 +- 错误处理 / 重试逻辑 + +## 按语言的搜索模式 + +### Go +- 接口: `type \w+ interface` +- 结构体: `type \w+ struct` +- 函数签名: `func \([^)]+\) \w+` +- 依赖注入: `New\w+\(.*\w+Client` +- Goroutine/Channel: `go func`, `chan ` +- 错误处理: `if err != nil` + +### Python +- 类: `class \w+` +- 函数: `def \w+` +- 装饰器: `@\w+`(路由、依赖注入) +- 异步: `async def`, `await ` +- 导入: `from .* import`, `import ` + +### TypeScript/JavaScript +- 类: `class \w+` +- 接口: `interface \w+` +- 导入: `import .* from` +- 路由: `app\.(get|post|put|delete)` +- 中间件: `\.use\(` + +## 信息提取深度 + +- **架构图/组件图/部署图**: 只需包级/模块级信息,读 CLAUDE.md + 入口文件即可 +- **序列图/通信图**: 需要函数调用链,读关键源码文件 +- **类图/对象图**: 需要类型定义,读 types/model 文件 +- **流程图/活动图/泳道图**: 需要主流程代码,读 pipeline/orchestrator/handler 文件 +- **数据流图**: 需要数据结构 + 变换逻辑,读 processor/converter 文件 +- **用例图**: 读 CLAUDE.md + router/api 文件 diff --git a/pigo/internal/builtinskills/skills/insight-diagram/scripts/review_svg.py b/pigo/internal/builtinskills/skills/insight-diagram/scripts/review_svg.py new file mode 100644 index 0000000..6a9a4fc --- /dev/null +++ b/pigo/internal/builtinskills/skills/insight-diagram/scripts/review_svg.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 +"""审查 insight-diagram 生成的 SVG 图,做几何校验。 + +检查项(对应三条要求): + 1. 箭头两端是否落在框图边缘线上 —— 既不"深入"框内,也不"空接"悬空。 + 2. 非嵌套框图之间是否重叠。 + 3. 框图之间是否留出足够间距。 + +用法: + python3 review_svg.py docs/architecture.html [more.html ...] + python3 review_svg.py docs/*.html --min-gap 8 --json + +退出码: 发现 ERROR 时为 1;仅 WARNING(或干净)为 0;加 --strict 让 WARNING 也返回 1。 + +实现说明:用 html.parser(而非 XML 解析器)遍历,以兼容 SVG-in-HTML 中 +未转义的 & 或文本里的 <> 等;箭头端点允许落在任意图形边缘、 +任意连线/生命线上(汇合点),故序列图/通信图等不会误报"空接"。 +""" +import argparse +import json +import math +import re +import sys +from html.parser import HTMLParser + +# ---- 几何容差(像素)---- +BOUNDARY_TOL = 6.0 # 端点距图形边 <= 此值视为"落在边上",正常 +PENETRATION = 8.0 # 端点在框内且距最近边 > 此值视为"深入"(ERROR) +FLOATING = 8.0 # 端点距所有图形 > 此值且不接任何连线 → "空接"(WARNING) +JUNCTION_TOL = 6.0 # 端点距其他连线/生命线 <= 此值视为合法汇合点 +OVERLAP_EPS = 2.0 # 两方向交叠都 > 此值才算重叠 +CONTAIN_MARGIN = 2.0 # 判定包含时允许内框略微超出 +SAME_BOX_TOL = 3.0 # 各边相差 <= 此值视为同一个框(去重) + +# 只有 w/h 同时达到阈值的图形才算"框图节点",参与重叠/间距校验; +# 更小的(标签底衬、终止圆点、图例色块)仅作为箭头落点目标。 +MIN_BOX_W = 36.0 +MIN_BOX_H = 22.0 + +# 这些子树内的图形只是装饰/定义,不收集 +SKIP_SUBTREES = {'defs', 'marker', 'pattern', 'clippath', 'lineargradient', + 'radialgradient', 'symbol', 'mask'} + + +# ===================================================================== +# 基础工具 +# ===================================================================== +def _floats(s): + return [float(x) for x in re.findall(r'-?\d+(?:\.\d+)?', s or '')] + + +def _num(v): + """解析 SVG 数值属性;含 '%' 或无法解析时返回 None。""" + if v is None or '%' in v: + return None + m = re.match(r'\s*(-?\d+(?:\.\d+)?)', v) + return float(m.group(1)) if m else None + + +def _parse_translate(transform): + """累加 transform 中的 translate 偏移,返回 (dx, dy)。""" + dx = dy = 0.0 + for m in re.finditer(r'translate\(\s*([-\d.]+)[\s,]*([-\d.]+)?\s*\)', transform or ''): + dx += float(m.group(1)) + dy += float(m.group(2)) if m.group(2) is not None else 0.0 + return dx, dy + + +def _point_seg_dist(px, py, ax, ay, bx, by): + """点到线段的最短距离。""" + dx, dy = bx - ax, by - ay + if dx == 0 and dy == 0: + return math.hypot(px - ax, py - ay) + t = ((px - ax) * dx + (py - ay) * dy) / (dx * dx + dy * dy) + t = max(0.0, min(1.0, t)) + return math.hypot(px - (ax + t * dx), py - (ay + t * dy)) + + +def parse_path_points(d): + """提取 path 各命令的落点(曲线只取终点),用于求首/末锚点。""" + pts = [] + cx = cy = 0.0 + cmd, nums = None, [] + tokens = re.findall(r'([MmLlHhVvCcSsQqTtAaZz])|(-?\d+(?:\.\d+)?)', d or '') + + def flush(): + nonlocal cx, cy + if cmd in ('M', 'L', 'T'): + for i in range(0, len(nums) - 1, 2): + cx, cy = nums[i], nums[i + 1]; pts.append((cx, cy)) + elif cmd in ('m', 'l', 't'): + for i in range(0, len(nums) - 1, 2): + cx, cy = cx + nums[i], cy + nums[i + 1]; pts.append((cx, cy)) + elif cmd == 'H': + for v in nums: cx = v; pts.append((cx, cy)) + elif cmd == 'h': + for v in nums: cx += v; pts.append((cx, cy)) + elif cmd == 'V': + for v in nums: cy = v; pts.append((cx, cy)) + elif cmd == 'v': + for v in nums: cy += v; pts.append((cx, cy)) + elif cmd in ('C', 'S', 'Q') and len(nums) >= 2: + cx, cy = nums[-2], nums[-1]; pts.append((cx, cy)) + elif cmd in ('c', 's', 'q') and len(nums) >= 2: + cx, cy = cx + nums[-2], cy + nums[-1]; pts.append((cx, cy)) + for tok_cmd, tok_num in tokens: + if tok_cmd: + if cmd is not None: + flush() + cmd, nums = tok_cmd, [] + else: + nums.append(float(tok_num)) + if cmd is not None: + flush() + return pts + + +# ===================================================================== +# 图形对象 +# ===================================================================== +class Shape: + """任意几何图形的统一表示,用 bbox + 类型描述。""" + __slots__ = ('kind', 'x0', 'y0', 'x1', 'y1', 'cx', 'cy', 'rx', 'ry') + + def __init__(self, kind, x0, y0, x1, y1): + self.kind = kind + self.x0, self.y0, self.x1, self.y1 = x0, y0, x1, y1 + self.cx, self.cy = (x0 + x1) / 2, (y0 + y1) / 2 + self.rx, self.ry = (x1 - x0) / 2, (y1 - y0) / 2 + + @property + def w(self): return self.x1 - self.x0 + @property + def h(self): return self.y1 - self.y0 + + def is_box(self): + return self.w >= MIN_BOX_W and self.h >= MIN_BOX_H + + def signed_dist(self, px, py): + """点到边界的有符号距离:内部为负、外部为正、≈0 在边上。""" + if self.kind == 'ellipse' and self.rx > 0 and self.ry > 0: + nx, ny = (px - self.cx) / self.rx, (py - self.cy) / self.ry + return (math.hypot(nx, ny) - 1.0) * ((self.rx + self.ry) / 2.0) + if self.kind == 'circle' and self.rx > 0: + return math.hypot(px - self.cx, py - self.cy) - self.rx + dx = max(self.x0 - px, 0, px - self.x1) + dy = max(self.y0 - py, 0, py - self.y1) + if dx > 0 or dy > 0: + return math.hypot(dx, dy) + return -min(px - self.x0, self.x1 - px, py - self.y0, self.y1 - py) + + def bbox_key(self): + return (round(self.x0, 1), round(self.y0, 1), + round(self.x1, 1), round(self.y1, 1)) + + +def _contains(a, b): + """框 a 是否(在容差内)包含框 b 且二者不等大。""" + return (b.x0 >= a.x0 - CONTAIN_MARGIN and b.x1 <= a.x1 + CONTAIN_MARGIN and + b.y0 >= a.y0 - CONTAIN_MARGIN and b.y1 <= a.y1 + CONTAIN_MARGIN and + not (abs(a.x0 - b.x0) < SAME_BOX_TOL and abs(a.x1 - b.x1) < SAME_BOX_TOL and + abs(a.y0 - b.y0) < SAME_BOX_TOL and abs(a.y1 - b.y1) < SAME_BOX_TOL)) + + +# ===================================================================== +# 用 HTMLParser 遍历 SVG(容忍未转义字符) +# ===================================================================== +class SvgCollector(HTMLParser): + def __init__(self): + super().__init__(convert_charrefs=True) + self.stack = [{'ox': 0.0, 'oy': 0.0, 'skip': False}] + self.shapes = [] # 所有图形(含小图形),用于箭头落点目标 + self.connectors = [] # 带 marker 的 line/path + self.segments = [] # 所有 line/path 折线段,用于汇合点判定 + + # void/自闭合元素 + def handle_startendtag(self, tag, attrs): + self._emit(tag.lower(), dict(attrs)) + + def handle_starttag(self, tag, attrs): + tag = tag.lower() + a = dict(attrs) + parent = self.stack[-1] + dx, dy = _parse_translate(a.get('transform', '')) + node = {'ox': parent['ox'] + dx, 'oy': parent['oy'] + dy, + 'skip': parent['skip'] or tag in SKIP_SUBTREES} + self.stack.append(node) + self._emit(tag, a, ctx=node) + + def handle_endtag(self, tag): + if len(self.stack) > 1: + self.stack.pop() + + def _emit(self, tag, a, ctx=None): + ctx = ctx or self.stack[-1] + if ctx['skip'] or tag in SKIP_SUBTREES: + return + ox, oy = ctx['ox'], ctx['oy'] + dx, dy = _parse_translate(a.get('transform', '')) # 自闭合元素自身的 translate + if tag in ('rect', 'ellipse', 'circle', 'polygon'): + ox, oy = ox + dx, oy + dy + + if tag == 'rect': + x, y = _num(a.get('x', '0')), _num(a.get('y', '0')) + w, h = _num(a.get('width')), _num(a.get('height')) + if None not in (x, y, w, h): + self.shapes.append(Shape('rect', ox + x, oy + y, ox + x + w, oy + y + h)) + elif tag == 'ellipse': + cx, cy = _num(a.get('cx', '0')), _num(a.get('cy', '0')) + rx, ry = _num(a.get('rx')), _num(a.get('ry')) + if None not in (cx, cy, rx, ry): + self.shapes.append(Shape('ellipse', ox + cx - rx, oy + cy - ry, + ox + cx + rx, oy + cy + ry)) + elif tag == 'circle': + cx, cy = _num(a.get('cx', '0')), _num(a.get('cy', '0')) + r = _num(a.get('r')) + if None not in (cx, cy, r): + self.shapes.append(Shape('circle', ox + cx - r, oy + cy - r, + ox + cx + r, oy + cy + r)) + elif tag == 'polygon': + nums = _floats(a.get('points', '')) + pts = list(zip(nums[0::2], nums[1::2])) + if len(pts) >= 3: + xs = [ox + p[0] for p in pts]; ys = [oy + p[1] for p in pts] + self.shapes.append(Shape('polygon', min(xs), min(ys), max(xs), max(ys))) + elif tag == 'line': + x1, y1 = _num(a.get('x1', '0')), _num(a.get('y1', '0')) + x2, y2 = _num(a.get('x2', '0')), _num(a.get('y2', '0')) + if None not in (x1, y1, x2, y2): + seg = [(ox + x1, oy + y1), (ox + x2, oy + y2)] + self.segments.append(seg) + if a.get('marker-end') or a.get('marker-start'): + self.connectors.append({ + 'a': seg[0], 'b': seg[-1], 'seg': seg, + 'arrow_a': bool(a.get('marker-start')), + 'arrow_b': bool(a.get('marker-end'))}) + elif tag == 'path': + pts = [(ox + px, oy + py) for px, py in parse_path_points(a.get('d', ''))] + if len(pts) >= 2: + self.segments.append(pts) + if a.get('marker-end') or a.get('marker-start'): + self.connectors.append({ + 'a': pts[0], 'b': pts[-1], 'seg': pts, + 'arrow_a': bool(a.get('marker-start')), + 'arrow_b': bool(a.get('marker-end'))}) + + +def collect(svg_text): + """返回 (boxes, all_shapes, connectors, segments)。""" + p = SvgCollector() + p.feed(svg_text) + # 框去重(描边 + 遮罩底衬常画两层完全重合的 rect) + seen, boxes = set(), [] + for s in p.shapes: + if s.is_box(): + k = s.bbox_key() + if k not in seen: + seen.add(k) + boxes.append(s) + return boxes, p.shapes, p.connectors, p.segments + + +def extract_svg(text): + m = re.search(r'', text, re.DOTALL | re.IGNORECASE) + return m.group(0) if m else None + + +# ===================================================================== +# 三项检查 +# ===================================================================== +def check_arrow_endpoints(shapes, connectors, segments): + """检查 1:箭头端点应恰好落在某图形边缘,或合法汇入另一连线。 + + 会先剔除"装饰性"连线:两端都既不贴任何图形边、也不汇入其它线段 + (典型如图例 Legend 里的示例箭头 / 独立标注线),不参与判定。 + """ + issues = [] + + def status(px, py): + dists = [s.signed_dist(px, py) for s in shapes] + on_edge = bool(dists) and any(abs(d) <= BOUNDARY_TOL for d in dists) + deepest = min(dists) if dists else 0.0 + nearest_out = min((d for d in dists if d >= 0), default=None) + return on_edge, deepest, nearest_out + + for i, c in enumerate(connectors): + sa = status(*c['a']) + sb = status(*c['b']) + a_anchored = sa[0] or _near_other_segment(*c['a'], segments, c['seg']) + b_anchored = sb[0] or _near_other_segment(*c['b'], segments, c['seg']) + # 两端都不锚定 → 视为图例/装饰线,跳过 + if not a_anchored and not b_anchored: + continue + + ends = [] + if c['arrow_b']: + ends.append(('终点', c['b'], sb)) + if c['arrow_a']: + ends.append(('起点', c['a'], sa)) + for label, (px, py), (on_edge, deepest, nearest_out) in ends: + if on_edge: + continue # 落在某图形边上:正常 + if deepest < -PENETRATION: + issues.append(('ERROR', + f'连线#{i+1} {label}({px:.0f},{py:.0f}) 深入框内 ' + f'{-deepest:.0f}px,应止于框边缘')) + continue + if nearest_out is not None and nearest_out > FLOATING: + if _near_other_segment(px, py, segments, c['seg']): + continue # 汇入另一连线/生命线 + issues.append(('WARNING', + f'连线#{i+1} {label}({px:.0f},{py:.0f}) 悬空,' + f'距最近框边 {nearest_out:.0f}px(空接)')) + return issues + + +def _near_other_segment(px, py, segments, own): + for seg in segments: + if seg is own: + continue + for k in range(len(seg) - 1): + if _point_seg_dist(px, py, *seg[k], *seg[k + 1]) <= JUNCTION_TOL: + return True + return False + + +def check_overlap(boxes): + """检查 2:非嵌套框之间不得重叠。""" + issues = [] + for i in range(len(boxes)): + for j in range(i + 1, len(boxes)): + a, b = boxes[i], boxes[j] + if _contains(a, b) or _contains(b, a): + continue + ox = min(a.x1, b.x1) - max(a.x0, b.x0) + oy = min(a.y1, b.y1) - max(a.y0, b.y0) + if ox > OVERLAP_EPS and oy > OVERLAP_EPS: + issues.append(('ERROR', + f'框[{a.x0:.0f},{a.y0:.0f} {a.w:.0f}x{a.h:.0f}] 与 ' + f'[{b.x0:.0f},{b.y0:.0f} {b.w:.0f}x{b.h:.0f}] ' + f'重叠 {ox:.0f}x{oy:.0f}px')) + return issues + + +def check_spacing(boxes, min_gap): + """检查 3:投影相邻、不嵌套、不重叠的框,净间距需 >= min_gap。""" + issues = [] + for i in range(len(boxes)): + for j in range(i + 1, len(boxes)): + a, b = boxes[i], boxes[j] + if _contains(a, b) or _contains(b, a): + continue + xo = min(a.x1, b.x1) - max(a.x0, b.x0) + yo = min(a.y1, b.y1) - max(a.y0, b.y0) + if xo > OVERLAP_EPS and yo > OVERLAP_EPS: + continue # 重叠交给 check_overlap + gap, axis = None, '' + if xo > OVERLAP_EPS: + gap, axis = max(a.y0, b.y0) - min(a.y1, b.y1), '垂直' + elif yo > OVERLAP_EPS: + gap, axis = max(a.x0, b.x0) - min(a.x1, b.x1), '水平' + if gap is not None and 0 <= gap < min_gap: + issues.append(('WARNING', + f'框[{a.x0:.0f},{a.y0:.0f}] 与 [{b.x0:.0f},{b.y0:.0f}] ' + f'{axis}间距仅 {gap:.0f}px (< {min_gap:.0f}px)')) + return issues + + +def review_file(path, min_gap): + try: + with open(path, encoding='utf-8') as f: + text = f.read() + except OSError as e: + return {'file': path, 'issues': [('ERROR', f'无法读取: {e}')], + 'boxes': 0, 'connectors': 0} + svg = extract_svg(text) + if not svg: + return {'file': path, 'issues': [('ERROR', '未找到 块')], + 'boxes': 0, 'connectors': 0} + boxes, shapes, connectors, segments = collect(svg) + issues = (check_arrow_endpoints(shapes, connectors, segments) + + check_overlap(boxes) + + check_spacing(boxes, min_gap)) + return {'file': path, 'issues': issues, + 'boxes': len(boxes), 'connectors': len(connectors)} + + +# ===================================================================== +# CLI +# ===================================================================== +def main(argv=None): + ap = argparse.ArgumentParser( + description='审查 insight-diagram 生成的 SVG(箭头落点 / 框重叠 / 框间距)') + ap.add_argument('files', nargs='+', help='待检查的 HTML/SVG 文件') + ap.add_argument('--min-gap', type=float, default=8.0, + help='相邻框最小净间距阈值 px,默认 8') + ap.add_argument('--json', action='store_true', help='以 JSON 输出') + ap.add_argument('--strict', action='store_true', + help='存在 WARNING 时也以非零码退出') + args = ap.parse_args(argv) + + results = [review_file(p, args.min_gap) for p in args.files] + + if args.json: + print(json.dumps([ + {'file': r['file'], 'boxes': r['boxes'], 'connectors': r['connectors'], + 'issues': [{'level': lv, 'message': m} for lv, m in r['issues']]} + for r in results], ensure_ascii=False, indent=2)) + else: + for r in results: + errs = [m for lv, m in r['issues'] if lv == 'ERROR'] + warns = [m for lv, m in r['issues'] if lv == 'WARNING'] + mark = '✗' if errs else ('⚠' if warns else '✓') + print(f'\n{mark} {r["file"]} ' + f'({r["boxes"]} 框 / {r["connectors"]} 箭头连线)') + for m in errs: + print(f' ERROR {m}') + for m in warns: + print(f' WARNING {m}') + if not errs and not warns: + print(' 通过:箭头落点、框重叠、框间距均无异常') + + total_err = sum(1 for r in results for lv, _ in r['issues'] if lv == 'ERROR') + total_warn = sum(1 for r in results for lv, _ in r['issues'] if lv == 'WARNING') + if not args.json: + print(f'\n汇总:{total_err} 个 ERROR,{total_warn} 个 WARNING,' + f'共 {len(results)} 个文件') + return 1 if (total_err or (args.strict and total_warn)) else 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/pigo/internal/builtinskills/skills/loop-it/SKILL.md b/pigo/internal/builtinskills/skills/loop-it/SKILL.md new file mode 100644 index 0000000..0b9bd83 --- /dev/null +++ b/pigo/internal/builtinskills/skills/loop-it/SKILL.md @@ -0,0 +1,562 @@ +--- +name: loop-it +description: "Automated issue loop with checkpoint/resume: fetch open GitHub issues → dependency-aware topological sort → implement each issue end-to-end → review with /review-it → document with /note-it → ship with /ship-it → repeat. Persists state to .loop-state.json for crash recovery. Triggers on: loop-it, loop issues, auto implement, 批量实现, 循环实现, 实现所有issue, 恢复循环, resume loop." +user-invocable: true +allowed-tools: + - Bash(gh:*) + - Bash(git:*) + - Bash(cat:*) +--- + +# loop-it — 带检查点恢复的自动化 Issue 实现循环 + +Fetch all open GitHub issues, resolve dependency order, implement each through the full pipeline (内联实现 → `/review-it` → `/note-it` → `/ship-it`), persist progress to state file, and resume from checkpoint on crash. + +> **⚠️ 关键前提:实现步骤由 agent 内联自主完成,不依赖任何外部 `/goal` 命令。** +> 本环境中不存在可调用的 `goal` 命令或 skill。因此「实现 issue」这一步**必须由 agent 内联完成**:直接读取该 issue 的标题与正文(含其引用的 PRD/SPEC 与验收条件),自主完成"理解需求 → 写/改代码 → 跑测试与 lint → 满足全部验收条件"的闭环,持续工作直到该 issue 的验收条件全部满足且测试/构建通过。**不要**尝试用 Skill 工具调用 `goal`(会报 `goal is a UI command, not a skill`),也**不要**因为找不到 `/goal` 而中止循环。`/review-it`、`/note-it`、`/ship-it` 仍是真实 skill,经 Skill 工具调用。 + +--- + +## Overview + +``` +前置检查 → 读取状态文件 → Fetch Issues → 构建依赖图 → 拓扑排序 + | + ┌───────────────────────────────────────────────────────────┘ + | + v +┌──────────────── 单 Issue 循环 ────────────────┐ +| | +| 从检查点恢复?—— 跳过已完成/失败的 | +| | +| 分支准备 (checkout main, pull, create branch) | +| | | +| Skip/Blocked? ── 是 → 标记 skipped/blocked, 写检查点 | +| | | +| 否 | +| | | +| 内联实现 → 出错?→ 分类 → 恢复 → 重试 | +| | | | +| | 失败 → 检查点, 下一个 | +| | | +| /review-it → 有问题?→ 修复 → 重跑 review | +| | | +| /note-it (捕获实现笔记, best-effort) | +| | | +| /ship-it → 出错?→ 分类 → 恢复 | +| | | +| 分支清理 (checkout main, pull, delete branch) | +| | | +| 检查点 (标记 shipped) | +| | | +└────────┴──────────────────────────────────────┘ + | + v + 全部完成 → 最终 Summary +``` + +--- + +## 前置检查 + +开始循环前,按顺序验证所有前提条件。任何检查失败则停止并打印错误。 + +### Check 1: gh CLI 认证 + +```bash +gh auth status +``` + +失败 → 打印 `❌ gh CLI 未认证。运行: gh auth login`,退出。 + +### Check 2: Git 仓库 + +```bash +git rev-parse --is-inside-work-tree +``` + +失败 → 打印 `❌ 不在 git 仓库中`,退出。 + +### Check 3: Git 工作树清洁度 + +```bash +git status --porcelain +``` + +有输出(dirty)→ 打印 `⚠️ 工作树有未提交的更改`,提供选项: +- A. `git stash` 暂存后继续 +- B. 中止,让用户自行处理 +- C. 强制继续(不推荐) + +默认 B。 + +### Check 4: 在默认分支上 + +```bash +git branch --show-current +``` + +不在 main/master → 打印 `⚠️ 当前在 {branch} 分支`,提供选项: +- A. `git checkout main && git pull` 切换 +- B. 继续在当前分支 + +### Check 5: 远程可达 + +```bash +git ls-remote --heads origin +``` + +失败 → 打印 `❌ 无法访问远程仓库。检查网络和权限`,退出。 + +### Check 6: 状态文件存在? + +```bash +cat .loop-state.json +``` + +存在 → 打印进度摘要,提供选项: +- A. 从检查点恢复 +- B. 从头开始(删除状态文件) +- C. 中止 + +--- + +## 状态文件 + +### 位置 + +`.loop-state.json`,放在 repo 根目录。**必须添加到 `.gitignore`**。如果文件被 git 跟踪,打印警告并建议用户添加到 `.gitignore`。 + +### 格式 + +```json +{ + "version": 1, + "started_at": "2025-06-09T10:00:00Z", + "updated_at": "2025-06-09T10:30:00Z", + "repo": "owner/repo-name", + "total_issues": 8, + "issues": { + "3": { + "status": "shipped", + "branch": "feat/issue-3-add-priority", + "started_at": "2025-06-09T10:00:00Z", + "completed_at": "2025-06-09T10:15:00Z", + "attempts": 1 + }, + "4": { + "status": "failed", + "phase": "goal", + "error_class": "build_failure", + "branch": "feat/issue-4-filter-tasks", + "started_at": "2025-06-09T10:15:00Z", + "updated_at": "2025-06-09T10:30:00Z", + "attempts": 3, + "last_error": "test TestFilterPriority failed: expected 3, got 0" + }, + "7": { + "status": "pending" + } + } +} +``` + +### 状态值 + +`pending` | `in_progress` | `skipped` | `shipped` | `failed` | `blocked` + +### 写入规则 + +- 每次状态转换后立即写入(`pending` → `in_progress`、`in_progress` → `shipped`/`failed`/`skipped` 等) +- 写入使用 `cat > .loop-state.json << 'LOOPSTATE'\n{json}\nLOOPSTATE` +- 如果状态文件已存在但内容损坏(非法 JSON),打印警告,提供从头开始或中止的选项。**绝不自动覆盖损坏文件** +- 循环完成后保留状态文件(作为记录),用户可手动删除 + +--- + +## Step 1: Fetch Issues & Build Dependency Graph + +Fetch all open issues: + +```bash +gh issue list --state open --json number,title,labels,body --limit 100 +``` + +### Parse Dependencies + +Read each issue body, look for patterns: +- `Dependencies: #3, #5` or `Depends on: #3` +- `depends on #3` or `requires #3` (in body text) + +Build a dependency graph. Sort using topological order: + +1. Issues with no dependencies first (sorted by number ascending) +2. Issues whose dependencies are all shipped/closed next +3. Blocked issues (depend on other open issues) last +4. Circular dependencies → print warning `⚠️ 循环依赖检测到: #A ↔ #B,按编号顺序处理`,break cycle by number order + +If no dependency patterns found in any issue body, fall back to number-ascending sort. + +Print ordered list: + +``` +📋 Found N open issues (topological sort): + #1: Add priority field (无依赖) + #3: Display indicator (依赖 #1) + #5: Add selector (依赖 #1) + #7: Filter view (依赖 #1, #3) +``` + +If no open issues → print `✅ No open issues found. Nothing to do.` and exit. + +--- + +## Step 2: Resume or Initialize + +### If `.loop-state.json` exists (from 前置检查 Check 6) + +1. Read the file +2. Print progress summary: + +``` +📊 从检查点恢复 (上次更新: {updated_at}) + ✅ Shipped: #1, #3 + ⏭️ Skipped: #2 (question) + ❌ Failed: #4 (build_failure — 3 attempts) + 📋 Remaining: #5, #7 +``` + +3. For each `failed` issue: ask user — retry or skip? +4. For `in_progress` issues: check if branch exists, changes exist → decide resume from current state or restart +5. Skip all `shipped`/`skipped` issues +6. Continue from first pending/retryable issue + +### If no state file + +1. Initialize new state file with all fetched issues as `pending` +2. Start from first issue in topological order + +--- + +## Step 3: Process Single Issue + +For each issue, print a banner: + +``` +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +🔄 Processing Issue #{number}: {title} +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +``` + +### 3a. Branch Prep + +Prepare a clean git environment for this issue: + +```bash +# 确认在 main 上 +git checkout main +git pull + +# 创建功能分支 +git checkout -b feat/issue-{N}-{short-desc} +``` + +Branch naming: `feat/issue-N-short-desc` or `fix/issue-N-short-desc`(与 /ship-it 保持一致) + +### 3b. Skip or Implement + +Read the issue title and body. Decide if it needs code implementation: + +**Skip if the issue is:** +- A question / discussion / clarification +- Documentation-only (typos, wording) +- Already implemented (check codebase) +- A duplicate of another issue +- Clearly labeled `wontfix`, `question`, `discussion`, or `invalid` +- Not actionable (no clear acceptance criteria and cannot infer any) + +**Skip (blocked) if the issue has unresolved dependencies:** +- Check the dependency graph from Step 1 +- If any dependency issue is not `shipped` (still `pending`, `failed`, `blocked`, or not in state file) → skip as blocked +- The dependency issue itself may have failed or been skipped — in either case, this issue cannot proceed safely + +When skipping: + +``` +⏭️ Skipping Issue #{number}: {title} + Reason: {why} +``` + +When blocked: + +``` +🔒 Blocking Issue #{number}: {title} + Reason: dependency #{dep_number} not shipped ({status}) +``` + +Update state: `pending` → `skipped` or `pending` → `blocked`, write checkpoint, run **3h Branch Cleanup**, proceed to next issue. + +### 3c. Implement (内联自主实现) + +Update state: `pending` → `in_progress`, `phase: "implement"`, write checkpoint. + +**由 agent 内联完成实现**(本环境无 `goal` 命令/skill 可调用,必须自己干): + +1. 读取该 issue 的标题与正文,提取需求与全部验收条件(Acceptance Criteria);若正文引用了 PRD/SPEC 文件(如 `tasks/prd-*.md`),一并读取作为上下文 +2. 阅读相关现有代码,遵循项目既有风格、命名与依赖约定 +3. 实现/修改代码以满足全部验收条件 +4. 跑项目的构建、测试与 lint(如 `go build ./...`、`go vet ./...`、`go test ./...`) +5. 持续工作直到该 issue 的验收条件**全部满足**且测试/构建/lint 通过 + +> 不要尝试用 Skill 工具调用 `goal`(会报 `goal is a UI command, not a skill`),也不要因找不到 `/goal` 而中止——实现就是你自己内联完成的工作。 + +**On success:** + +``` +✅ Issue #{number} implementation complete +``` + +Write checkpoint with `phase: "implement_done"`. + +**On failure** — classify error (see 错误分类与恢复), apply recovery strategy, retry up to max attempts. If all retries exhausted: + +``` +⚠️ Issue #{number} failed: {error_class} after {N} attempts + Manual intervention required. +``` + +Update state: `in_progress` → `failed`, write checkpoint, run **3h Branch Cleanup**, proceed to next issue. + +### 3d. Review with /review-it + +Write checkpoint with `phase: "review"`. + +``` +/review-it +``` + +**If review finds actionable issues:** + +``` +🔍 Review found N issue(s) for #{number}. Fixing... +``` + +Fix each accepted finding, re-run `/review-it`. Repeat until clean or max 2 review rounds. + +**If review is clean:** + +``` +✅ Review clean for Issue #{number} +``` + +Write checkpoint with `phase: "review_done"`. + +### 3e. Document with /note-it + +After review, before ship — capture implementation notes: + +``` +/note-it +``` + +This creates `docs/issue#XXXX.html` with design decisions, deviations, tradeoffs, and open questions. + +**On success:** + +``` +📝 Issue #{number} notes captured +``` + +**On failure** (can't determine issue number, etc.) — print warning but **do not block shipping**: + +``` +⚠️ /note-it failed for Issue #{number}: {reason}. Continuing to ship. +``` + +Write checkpoint with `phase: "note_done"`. + +### 3f. Ship with /ship-it + +``` +/ship-it +``` + +This commits, pushes, creates PR, merges, and closes the issue. + +**On success:** + +``` +🚀 Issue #{number} shipped successfully! +``` + +**On failure** — classify error (see 错误分类与恢复), apply recovery. If unresolvable: + +``` +⚠️ Issue #{number} ship failed: {error}. Manual merge required. +``` + +Update state: `in_progress` → `failed`, `phase: "ship"`, write checkpoint, run **3h Branch Cleanup**, proceed to next issue. + +### 3g. Checkpoint + +After successful ship, update state: `in_progress` → `shipped`, set `completed_at`, write checkpoint. + +Print progress (see 进度可观测性). + +### 3h. Branch Cleanup + +After each issue (shipped, skipped, or failed): + +```bash +# 切回 main +git checkout main +git pull + +# 删除本地功能分支(仅当 shipped 时) +git branch -d feat/issue-{N}-{short-desc} +``` + +**For failed issues**: do NOT delete the branch. Keep it for investigation. + +### Next Issue + +Return to Step 3 for the next issue in topological order. + +When all issues processed, print final summary: + +``` +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +📊 Loop Complete — Summary +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + ✅ Shipped: N issues (#1, #3, ...) + ⏭️ Skipped: N issues (#2 — reason, #5 — reason, ...) + 🔒 Blocked: N issues (#7 — depends on #4, ...) + ❌ Failed: N issues (#4 — error, ...) + 📋 Total: N issues +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +``` + +--- + +## 错误分类与恢复 + +当错误发生时,先分类,再按策略恢复。 + +| 错误类别 | 检测信号 | 恢复策略 | 最大重试 | +|----------|---------|---------|---------| +| build_failure | 编译错误、undefined、类型错误 | 读错误,修代码,重新构建 | 3 | +| test_failure | 断言失败、test failed | 读测试输出,修实现,重跑测试 | 3 | +| lint_failure | lint 错误、格式问题 | 自动修复 (lint --fix),重跑 | 2 | +| merge_conflict | CONFLICT 标记 | rebase origin/main,解决冲突,push | 2 | +| ci_failure | gh pr checks 失败 | 读 CI 日志,本地修复,push | 2 | +| auth_failure | 403、401、认证错误 | 停止,告知用户重新认证 | 0 | +| rate_limit | rate limit、secondary abuse | 等待 60s,重试 | 3 | +| issue_unclear | issue 无验收条件且无法推断需求 | 跳过,标记 failed | 0 | +| network_error | timeout、connection refused | 等待 30s,重试 | 3 | +| unknown | 其他情况 | 记录完整错误,跳过 | 0 | + +**恢复协议:** + +1. 匹配错误类别 +2. 匹配成功 → 应用恢复策略,重试最多 N 次 +3. 重试全部失败 → 标记 `failed`,写检查点,继续下一个 issue +4. 无法匹配 → 标记 `failed`(error_class: `unknown`),继续 +5. **绝不无限重试。绝不未经确认 force-push。** + +--- + +## 进度可观测性 + +每完成一个 issue 后,打印结构化进度: + +``` +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +📊 Progress: 3/8 issues (37%) +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + ✅ Shipped: #1, #3 + ⏭️ Skipped: #2 (question), #6 (duplicate) + 🔒 Blocked: #7 (depends on #4 — failed) + ❌ Failed: #4 (build_failure — 3 attempts) + 📋 Remaining: #5, #8 +━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +``` + +--- + +## Logging Rules + +Every key step MUST print a log line with emoji prefix: + +| Emoji | Meaning | +|-------|---------| +| 📋 | Fetch / list | +| 🔄 | Processing issue | +| ⏭️ | Skip | +| 🔒 | Blocked (dependency not shipped) | +| ✅ | Success | +| ❌ | Failure | +| 🔍 | Review | +| 📝 | Notes (/note-it) | +| 🚀 | Ship | +| ⚠️ | Warning / retry | +| 📊 | Progress / summary | + +--- + +## Safety Guards + +- **Never force-push to main/master** — always use feature branches +- **Never skip review** — always run `/review-it` before `/ship-it` +- **Never skip notes** — always run `/note-it` before `/ship-it`(best-effort,不阻塞) +- **Max retries per error class** — 参见错误分类与恢复表,不无限重试 +- **Max 2 review rounds** — don't over-polish +- **Pause on CI failure** — log and continue, don't auto-override branch protection +- **Preserve issue labels** — only close issues that were actually shipped +- **Never auto-delete failed branches** — 保留供调查 +- **Checkpoint at every transition** — 每次状态变更写检查点,不仅仅在 ship 时 +- **State file integrity** — 损坏时警告用户,绝不自动覆盖 +- **State file in .gitignore** — 提醒用户添加 `.loop-state.json` +- **Strictly sequential** — 一次只处理一个 issue(实现会修改工作树,不能并行) +- **Skip dependency-blocked issues** — 依赖的 issue 未 shipped 时标记 `blocked` +- **实现由 agent 内联完成** — 本环境无 `goal` 命令/skill 可调用;「实现 issue」必须由 agent 自己读 issue、写代码、跑测试完成。报 `goal is a UI command, not a skill` 时不要中止,直接内联实现 + +--- + +## Edge Cases + +| Scenario | Handling | +|----------|----------| +| No open issues | Print "nothing to do" and exit | +| All issues are questions | Skip all, report summary | +| `gh` not authenticated | Print error, suggest `gh auth login`, exit | +| Issue has no body | Use title only to decide skip/implement | +| Issue references PRD/SPEC | 读取被引用的 PRD/SPEC 作为上下文,agent 内联实现 | +| Multiple issues depend on each other | Topological sort; dependencies already shipped first | +| Git repo is dirty before starting | 前置检查 Check 3: stash/abort/force | +| State file corrupted (invalid JSON) | 警告用户,提供从头开始或中止选项。绝不自动覆盖 | +| State file from different repo | 检测 repo 字段不匹配,警告,提供从头开始选项 | +| Issue `in_progress` from previous run | 检查分支是否存在、是否有变更 → 恢复或重新开始 | +| User aborts mid-loop | 状态文件已包含最新检查点,下次运行可恢复 | +| New issues created during loop | 不重新获取。完成当前批次后运行新 `/loop-it` | +| Circular dependencies | 打印警告,按编号顺序打破循环 | +| `/note-it` can't find issue number | 打印警告,跳过 /note-it,继续 /ship-it | +| `.loop-state.json` is git-tracked | 警告用户添加到 .gitignore,继续 | +| 误以为需要外部 `goal` 命令 | 本环境无此命令;「实现 issue」由 agent 内联完成(读 issue → 写代码 → 测试),不要中止循环 | + +--- + +## Relationship to Other Skills + +``` +/loop-it + ├── 内联实现 ← implement each issue(agent 自主读 issue、写代码、测试;非外部命令) + ├── /review-it ← review code before shipping(skill) + ├── /note-it ← capture implementation notes (best-effort)(skill) + └── /ship-it ← commit, PR, merge, close(skill) +``` + +Part of the goal-workflow pipeline: + +``` +/prd → /prd-to-spec → /to-issues → /loop-it (→ 内联实现 → /review-it → /note-it → /ship-it)×N +``` diff --git a/pigo/internal/builtinskills/skills/modern-go/SKILL.md b/pigo/internal/builtinskills/skills/modern-go/SKILL.md new file mode 100644 index 0000000..12303b7 --- /dev/null +++ b/pigo/internal/builtinskills/skills/modern-go/SKILL.md @@ -0,0 +1,1139 @@ +--- +name: modern-go +description: Modernize Go code by applying version-appropriate idioms and APIs (gofix-style transformations). Scans go.mod for the Go version, then transforms Go source files to use modern patterns—from Go 1.0 through 1.26+. Use when the user says "现代化","现代Go语言", "地道的", "idiomatic", "modernize", "modern-go", "update Go code", "gofix", or wants to upgrade Go idioms. +--- + +# modern-go + +Modernize Go source code by applying version-appropriate idioms, APIs, and language features. Works like `go fix` plus additional transformations curated from the Go team's modernize analysis passes and community best practices. + +## Usage + +Invoke this skill when the user asks to modernize Go code. By default, modernize the entire project; the user may specify a file or directory instead. + +When invoked: +1. Detect the project's Go version from `go.mod` (the `go` directive). +2. Find all `.go` files in the target scope (excluding `vendor/`, `.git/`, `testdata/`). +3. For each file, apply **all transformations for versions ≤ the project's Go version**, starting from the oldest to the newest. +4. After all transformations, print a summary of what was changed and what was skipped. + +If the user specifies a file or directory, limit the scope to that path. + +## Transformation Catalog + +Each transformation includes a **Go version** gate—only apply when the project's `go.mod` version ≥ that version. Never apply a transformation that requires a version higher than the project declares. + +### Go 1.0+ — `time.Since` + +| Before | After | +|---|---| +| `time.Now().Sub(start)` | `time.Since(start)` | + +```go +// before +elapsed := time.Now().Sub(start) +// after +elapsed := time.Since(start) +``` + +### Go 1.8+ — `time.Until` + +| Before | After | +|---|---| +| `deadline.Sub(time.Now())` | `time.Until(deadline)` | + +```go +// before +remaining := deadline.Sub(time.Now()) +// after +remaining := time.Until(deadline) +``` + +### Go 1.10+ — `strings.Builder` (loop concatenation) + +| Before | After | +|---|---| +| `s += item` in a loop | `var b strings.Builder; b.WriteString(item)` | + +```go +// before +s := "" +for _, item := range items { + s += item +} +// after +var b strings.Builder +for _, item := range items { + b.WriteString(item) +} +s := b.String() +``` + +Only when `+=` concatenation happens inside a loop. + +### Go 1.13+ — `errors.Is` + +| Before | After | +|---|---| +| `err == io.EOF` | `errors.Is(err, io.EOF)` | + +```go +// before +if err == io.EOF { + return +} +// after +if errors.Is(err, io.EOF) { + return +} +``` + +### Go 1.17+ — `//go:build` constraints (plusbuild) + +| Before | After | +|---|---| +| `// +build linux` + `//go:build linux` (both present) | keep only `//go:build linux` | + +```go +// before +//go:build linux && amd64 +// +build linux,amd64 + +package foo +// after +//go:build linux && amd64 + +package foo +``` + +The `plusbuild` modernizer removes obsolete `// +build` constraint lines once the equivalent `//go:build` line is present (the `//go:build` syntax landed in Go 1.17). Only strip the old line when a matching `//go:build` already exists — never drop the sole constraint. + +### Go 1.17+ — `unsafe.Add` / `unsafe.Slice` (unsafefuncs) + +| Before | After | +|---|---| +| `unsafe.Pointer(uintptr(ptr) + uintptr(n))` | `unsafe.Add(ptr, n)` | +| `(*[n]T)(unsafe.Pointer(p))[:]` slice construction | `unsafe.Slice(p, n)` | + +```go +// before — pointer arithmetic via uintptr +p2 := unsafe.Pointer(uintptr(ptr) + uintptr(offset)) +// after +p2 := unsafe.Add(ptr, offset) +``` + +```go +// before — building a slice from a base pointer +s := (*[1 << 30]byte)(unsafe.Pointer(p))[:n:n] +// after +s := unsafe.Slice(p, n) +``` + +The `unsafefuncs` modernizer (gopls v0.22.0) rewrites error-prone `uintptr` pointer math into `unsafe.Add` / `unsafe.Slice`, which the compiler and `go vet` understand as GC-safe. + +### Go 1.18+ — `any` + +| Before | After | +|---|---| +| `interface{}` | `any` | + +```go +// before +func decode(v interface{}) error { ... } +// after +func decode(v any) error { ... } +``` + +### Go 1.18+ — `strings.Cut` + +| Before | After | +|---|---| +| `i := strings.Index(s, sep); ... s[:i], s[i+len(sep):]` | `key, val, found := strings.Cut(s, sep)` | + +```go +// before +if i := strings.Index(s, "="); i >= 0 { + key, val := s[:i], s[i+1:] +} +// after +if key, val, found := strings.Cut(s, "="); found { + ... +} +``` + +### Go 1.18+ — `bytes.Cut` + +| Before | After | +|---|---| +| `i := bytes.Index(b, sep); ... b[:i], b[i+len(sep):]` | `before, after, found := bytes.Cut(b, sep)` | + +```go +// before +if i := bytes.Index(b, sep); i >= 0 { + before, after := b[:i], b[i+len(sep):] +} +// after +before, after, found := bytes.Cut(b, sep) +``` + +### Go 1.19+ — `fmt.Appendf` + +| Before | After | +|---|---| +| `buf = append(buf, fmt.Sprintf(...)...)` | `buf = fmt.Appendf(buf, ...)` | + +```go +// before +buf = append(buf, fmt.Sprintf("x=%d", x)...) +// after +buf = fmt.Appendf(buf, "x=%d", x) +``` + +### Go 1.19+ — Type-safe atomics (atomictypes) + +| Before | After | +|---|---| +| `atomic.StoreInt32(&v, 1)` / `atomic.LoadInt32(&v)` | `var v atomic.Int32; v.Store(1); v.Load()` | +| `atomic.AddInt64(&v, 1)` | `var v atomic.Int64; v.Add(1)` | +| `atomic.Value` + type assertion | `atomic.Pointer[T]` | + +```go +// before +var ready int32 +atomic.StoreInt32(&ready, 1) +if atomic.LoadInt32(&ready) == 1 { ... } + +// after +var ready atomic.Int32 +ready.Store(1) +if ready.Load() == 1 { ... } +``` + +```go +// before +var cache atomic.Value +cache.Store(&Config{}) +cfg := cache.Load().(*Config) + +// after +var cache atomic.Pointer[Config] +cache.Store(&Config{}) +cfg := cache.Load() +``` + +The `atomictypes` modernizer (gopls v0.22.0, `AtomicTypesAnalyzer`) rewrites both the variable declaration and every call site. Typed wrappers (`atomic.Int32/Int64/Uint32/Uint64/Bool/Pointer[T]`) have identical performance but prevent accidental non-atomic access and fix 64-bit alignment crashes on 32-bit architectures. + +### Go 1.20+ — `strings.Clone` + +| Before | After | +|---|---| +| `string([]byte(s))` | `strings.Clone(s)` | + +```go +// before +s2 := string([]byte(s)) // force copy +// after +s2 := strings.Clone(s) +``` + +### Go 1.20+ — `bytes.Clone` + +| Before | After | +|---|---| +| `make([]byte, len(src)); copy(dst, src)` | `bytes.Clone(src)` | + +```go +// before +dst := make([]byte, len(src)) +copy(dst, src) +// after +dst := bytes.Clone(src) +``` + +### Go 1.20+ — `strings.CutPrefix` / `strings.CutSuffix` + +| Before | After | +|---|---| +| `if strings.HasPrefix(s, p) { s = s[len(p):] }` | `if rest, ok := strings.CutPrefix(s, p); ok { s = rest }` | +| `if strings.HasSuffix(s, sf) { s = s[:len(s)-len(sf)] }` | `if rest, ok := strings.CutSuffix(s, sf); ok { s = rest }` | + +```go +// before +if strings.HasPrefix(s, "pre_") { + s = s[len("pre_"):] +} +// after +if rest, ok := strings.CutPrefix(s, "pre_"); ok { + s = rest +} +``` + +```go +// before +if strings.HasSuffix(s, ".txt") { + s = s[:len(s)-len(".txt")] +} +// after +if rest, ok := strings.CutSuffix(s, ".txt"); ok { + s = rest +} +``` + +### Go 1.20+ — `errors.Join` + +| Before | After | +|---|---| +| `fmt.Errorf("...: %w: %w", err1, err2)` | `errors.Join(err1, err2)` | + +```go +// before +return fmt.Errorf("load config: %w: %w", err1, err2) +// after +return errors.Join(fmt.Errorf("load config"), err1, err2) +``` + +### Go 1.20+ — `context.WithCancelCause` + +| Before | After | +|---|---| +| `ctx, cancel := context.WithCancel(parent)` + bare `cancel()` | `ctx, cancel := context.WithCancelCause(parent)` + `cancel(err)` | + +```go +// before +ctx, cancel := context.WithCancel(parent) +// ... somewhere ... +cancel() + +// after +ctx, cancel := context.WithCancelCause(parent) +cancel(ErrShutdown) +// caller: context.Cause(ctx) → ErrShutdown +``` + +### Go 1.21+ — `min` / `max` + +| Before | After | +|---|---| +| `if a < b { v = a } else { v = b }` | `v = min(a, b)` | +| `if a > b { v = a } else { v = b }` | `v = max(a, b)` | +| `if x < lo { x = lo }; if x > hi { x = hi }` | `x = min(max(x, lo), hi)` | + +```go +// before +lo := a +if b < lo { + lo = b +} +// after +lo := min(a, b) +``` + +```go +// before +if x < 0 { + x = 0 +} +if x > 100 { + x = 100 +} +// after +x = min(max(x, 0), 100) +``` + +### Go 1.21+ — `clear` + +| Before | After | +|---|---| +| `for k := range m { delete(m, k) }` | `clear(m)` | +| `for i := range s { s[i] = zero }` | `clear(s)` | + +```go +// before +for k := range m { + delete(m, k) +} +// after +clear(m) +``` + +```go +// before +for i := range s { + s[i] = 0 +} +// after +clear(s) +``` + +### Go 1.21+ — `slices` package + +| Before | After | +|---|---| +| Manual loop to find element | `slices.Contains(items, target)` | +| Loop returning index or -1 | `slices.Index(items, target)` | +| `sort.Slice(items, func(i,j int) bool { return items[i] < items[j] })` | `slices.SortFunc(items, cmp.Compare)` | +| Max/min finding loop | `slices.Max(items)` / `slices.Min(items)` | +| Reverse swap loop | `slices.Reverse(s)` | +| Remove consecutive duplicates loop | `slices.Compact(s)` | +| `s[:len(s):len(s)]` | `slices.Clip(s)` | +| `make([]T, len(src)); copy(dst, src)` | `slices.Clone(src)` | + +```go +// before → after: slices.Contains(items, target) +found := false +for _, x := range items { + if x == target { + found = true + break + } +} +``` + +```go +// before → after: slices.Index(items, target) +for i, x := range items { + if x == target { + return i + } +} +return -1 +``` + +```go +// before → after: slices.SortFunc(items, cmp.Compare) +sort.Slice(items, func(i, j int) bool { return items[i] < items[j] }) +``` + +```go +// before → after: slices.Max(items) / slices.Min(items) +max := items[0] +for _, v := range items[1:] { + if v > max { + max = v + } +} +``` + +```go +// before → after: slices.Reverse(s) +for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 { + s[i], s[j] = s[j], s[i] +} +``` + +```go +// before → after: slices.Compact(s) +i := 0 +for j := 1; j < len(s); j++ { + if s[j] != s[i] { + i++ + s[i] = s[j] + } +} +s = s[:i+1] +``` + +```go +// before → after: slices.Clip(s) +s = s[:len(s):len(s)] +``` + +```go +// before → after: slices.Clone(src) +dst := make([]T, len(src)) +copy(dst, src) +``` + +Requires importing `"slices"` and `"cmp"` (for `SortFunc`). + +### Go 1.21+ — `slices.Delete` / `slices.Insert` + +| Before | After | +|---|---| +| `append(s[:i], s[i+1:]...)` | `slices.Delete(s, i, i+1)` | +| `append(s[:i:i], append([]T{x}, s[i:]...)...)` | `slices.Insert(s, i, x)` | + +```go +// before — element removal (classic aliasing/leak footgun) +s = append(s[:i], s[i+1:]...) +// after +s = slices.Delete(s, i, i+1) +``` + +```go +// before — insert at index i +s = append(s[:i], append([]T{v}, s[i:]...)...) +// after +s = slices.Insert(s, i, v) +``` + +`slices.Delete` zeroes the tail elements to avoid retaining pointers (the manual `append` form leaks). Requires importing `"slices"`. + +### Go 1.21+ — `slices.Equal` / `maps.Equal` + +| Before | After | +|---|---| +| `reflect.DeepEqual(a, b)` for comparable slices | `slices.Equal(a, b)` | +| `reflect.DeepEqual(m1, m2)` for comparable maps | `maps.Equal(m1, m2)` | + +```go +// before +if reflect.DeepEqual(got, want) { ... } // got, want are []string +// after +if slices.Equal(got, want) { ... } +``` + +Faster and type-safe, with no reflection. Only for element types that are directly comparable (use `slices.EqualFunc` / `maps.EqualFunc` otherwise). Requires importing `"slices"` or `"maps"`. + +### Go 1.21+ — `maps` package + +| Before | After | +|---|---| +| Manual loop to copy a map | `maps.Clone(m)` | +| `for k, v := range src { dst[k] = v }` | `maps.Copy(dst, src)` | +| Loop + conditional delete | `maps.DeleteFunc(m, predicate)` | + +```go +// before → after: maps.Clone(m) +dst := make(map[K]V) +for k, v := range src { + dst[k] = v +} +``` + +```go +// before → after: maps.Copy(dst, src) +for k, v := range src { + dst[k] = v +} +``` + +```go +// before → after: maps.DeleteFunc(m, func(k K, v V) bool { return v == 0 }) +for k, v := range m { + if v == 0 { + delete(m, k) + } +} +``` + +Requires importing `"maps"`. + +### Go 1.22+ — `slices.Concat` (appendclipped) + +| Before | After | +|---|---| +| `append(append([]T(nil), s1...), s2...)` | `slices.Concat(s1, s2)` | +| `append(slices.Clip(s1), s2...)` for a fresh result | `slices.Concat(s1, s2)` | + +```go +// before +all := append(append([]int(nil), a...), b...) +// after +all := slices.Concat(a, b) +``` + +```go +// before — three-way concat +merged := append(append(append([]string(nil), x...), y...), z...) +// after +merged := slices.Concat(x, y, z) +``` + +The `appendclipped` modernizer replaces nested `append` concatenation of multiple slices with `slices.Concat`, which allocates a fresh, correctly-sized result. `slices.Concat` was added in Go 1.22. Requires importing `"slices"`. Only apply when the pattern builds a new slice (starts from `[]T(nil)` or a clipped base) — not when it appends in place to an existing slice. + +### Go 1.21+ — `sync.OnceFunc` / `sync.OnceValue` + +| Before | After | +|---|---| +| `var once sync.Once; once.Do(func() { ... })` | `f := sync.OnceFunc(func() { ... }); f()` | +| `sync.Once` + stored result variable | `sync.OnceValue(func() T { return val })` | + +```go +// before +var once sync.Once +func init() { once.Do(func() { setup() }) } + +// after +var initOnce = sync.OnceFunc(func() { setup() }) +``` + +```go +// before +var once sync.Once +var cfg *Config +func getConfig() *Config { + once.Do(func() { cfg = loadConfig() }) + return cfg +} +// after +var getConfig = sync.OnceValue(func() *Config { return loadConfig() }) +``` + +### Go 1.21+ — `context.AfterFunc` + +| Before | After | +|---|---| +| `go func() { <-ctx.Done(); cleanup() }()` | `stop := context.AfterFunc(ctx, cleanup)` | + +```go +// before +go func() { + <-ctx.Done() + conn.Close() +}() +// after +stop := context.AfterFunc(ctx, func() { conn.Close() }) +``` + +### Go 1.21+ — `context.WithTimeoutCause` / `WithDeadlineCause` + +| Before | After | +|---|---| +| `context.WithTimeout(parent, d)` | `context.WithTimeoutCause(parent, d, err)` | + +```go +// before +ctx, cancel := context.WithTimeout(parent, 5*time.Second) +// after +ctx, cancel := context.WithTimeoutCause(parent, 5*time.Second, ErrTimeout) +``` + +Only apply when a meaningful cause error is available. + +### Go 1.22+ — Range over integer + +| Before | After | +|---|---| +| `for i := 0; i < n; i++ { ... }` | `for i := range n { ... }` | +| `for i := 0; i < n; i++ { ... }` (i unused) | `for range n { ... }` | + +```go +// before +for i := 0; i < len(items); i++ { + process(i, items[i]) +} +// after +for i := range len(items) { + process(i, items[i]) +} +``` + +```go +// before +for i := 0; i < n; i++ { + doWork() +} +// after +for range n { + doWork() +} +``` + +### Go 1.22+ — Loop variable shadowing removal + +| Before | After | +|---|---| +| `for _, x := range items { x := x; ... }` | `for _, x := range items { ... }` | + +```go +// before +for _, x := range items { + x := x // capture for goroutine + go func() { use(x) }() +} +// after +for _, x := range items { + go func() { use(x) }() +} +``` + +The `x := x` capture idiom is redundant since Go 1.22. + +### Go 1.22+ — `cmp.Or` + +| Before | After | +|---|---| +| Chain of `if v == "" { v = fallback }` | `v := cmp.Or(val, fallback1, fallback2, ...)` | + +```go +// before +name := os.Getenv("NAME") +if name == "" { + name = os.Getenv("USER") +} +if name == "" { + name = "anonymous" +} +// after +name := cmp.Or(os.Getenv("NAME"), os.Getenv("USER"), "anonymous") +``` + +Requires importing `"cmp"`. + +### Go 1.22+ — `reflect.TypeFor` + +| Before | After | +|---|---| +| `reflect.TypeOf((*T)(nil)).Elem()` | `reflect.TypeFor[T]()` | + +```go +// before +t := reflect.TypeOf((*MyType)(nil)).Elem() +// after +t := reflect.TypeFor[MyType]() +``` + +### Go 1.22+ — Enhanced `http.ServeMux` + +| Before | After | +|---|---| +| `mux.HandleFunc("/api/", h)` + manual path parsing | `mux.HandleFunc("GET /api/{id}", h)` + `r.PathValue("id")` | + +```go +// before +mux.HandleFunc("/api/", func(w http.ResponseWriter, r *http.Request) { + id := strings.TrimPrefix(r.URL.Path, "/api/") + ... +}) +// after +mux.HandleFunc("GET /api/{id}", func(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + ... +}) +``` + +### Go 1.22+ — `math/rand/v2` + +| Before | After | +|---|---| +| `rand.Intn(n)` | `rand.IntN(n)` | +| `rand.Int63n(n)` / `rand.Int31n(n)` | `rand.Int64N(n)` / `rand.Int32N(n)` | +| `rand.Seed(...)` + global funcs | drop `Seed`; use auto-seeded top-level funcs or `rand.N[T]` | + +```go +// before +import "math/rand" +n := rand.Intn(100) +// after +import "math/rand/v2" +n := rand.IntN(100) +``` + +```go +// before — type-specific bound +d := time.Duration(rand.Int63n(int64(max))) +// after — generic N works for any integer type +d := rand.N(max) +``` + +`math/rand/v2` (Go 1.22) drops the deprecated global `Seed` (top-level funcs are auto-seeded) and adds generic `rand.N[T]`. **Semantic migration:** the random stream differs from `math/rand`, so do not apply where reproducibility from a fixed seed matters. Flag as a suggestion, not auto-apply. + +### Go 1.23+ — Range over function (iterators) + +| Before | After | +|---|---| +| Custom `func Walk(yield func(T) bool)` callback traversal | `func Walk() iter.Seq[T]` returning an iterator | + +```go +// before — callback-based iteration +func (t *Tree) Each(fn func(v int)) { + for _, v := range t.values { + fn(v) + } +} +// caller: t.Each(func(v int) { use(v) }) + +// after — standard iterator, usable with range +func (t *Tree) All() iter.Seq[int] { + return func(yield func(int) bool) { + for _, v := range t.values { + if !yield(v) { return } + } + } +} +// caller: for v := range t.All() { use(v) } +``` + +Adopt the `iter.Seq[T]` / `iter.Seq2[K,V]` protocol (Go 1.23) so custom containers compose with `range`, `slices.Collect`, `maps.Keys`, etc. Requires importing `"iter"`. Flag as a suggestion — it reshapes the API surface. + +### Go 1.23+ — Iterator helpers + +| Before | After | +|---|---| +| `var keys []K; for k := range m { keys = append(keys, k) }` | `slices.Collect(maps.Keys(m))` | +| `var vals []V; for _, v := range m { vals = append(vals, v) }` | `slices.Collect(maps.Values(m))` | + +```go +// before +var keys []string +for k := range m { + keys = append(keys, k) +} +// after +keys := slices.Collect(maps.Keys(m)) +``` + +```go +// before +var vals []int +for _, v := range m { + vals = append(vals, v) +} +// after +vals := slices.Collect(maps.Values(m)) +``` + +Requires importing `"slices"` and `"maps"`. + +### Go 1.23+ — `strings.SplitSeq` / `strings.FieldsSeq` + +| Before | After | +|---|---| +| `for _, part := range strings.Split(s, sep)` | `for part := range strings.SplitSeq(s, sep)` | +| `for _, field := range strings.Fields(s)` | `for field := range strings.FieldsSeq(s)` | + +```go +// before +for _, part := range strings.Split(line, ",") { + process(part) +} +// after +for part := range strings.SplitSeq(line, ",") { + process(part) +} +``` + +Only when the loop body does not need the index or the full slice. + +### Go 1.23+ — `bytes.SplitSeq` / `bytes.FieldsSeq` + +| Before | After | +|---|---| +| `for _, part := range bytes.Split(b, sep)` | `for part := range bytes.SplitSeq(b, sep)` | + +```go +// before +for _, part := range bytes.Split(data, sep) { + process(part) +} +// after +for part := range bytes.SplitSeq(data, sep) { + process(part) +} +``` + +### Go 1.23+ — `slices.Backward` (slicesbackward) + +| Before | After | +|---|---| +| `for i := len(s) - 1; i >= 0; i-- { use(s[i]) }` | `for _, v := range slices.Backward(s) { use(v) }` | +| reverse loop needing the index too | `for i, v := range slices.Backward(s) { ... }` | + +```go +// before +for i := len(items) - 1; i >= 0; i-- { + process(items[i]) +} +// after +for _, v := range slices.Backward(items) { + process(v) +} +``` + +```go +// before — index still needed +for i := len(items) - 1; i >= 0; i-- { + fmt.Println(i, items[i]) +} +// after +for i, v := range slices.Backward(items) { + fmt.Println(i, v) +} +``` + +The `slicesbackward` modernizer (gopls v0.22.0) replaces manual descending-index loops with the `slices.Backward` iterator. Requires importing `"slices"`. **Caveat:** the rewrite preserves exact semantics in normal cases, but do not apply it when the loop body mutates the slice length or the index is used for out-of-band arithmetic — those edge cases can become unsound. + +### Go 1.24+ — `t.Context()` in tests + +| Before | After | +|---|---| +| `ctx, cancel := context.WithCancel(context.Background()); defer cancel()` | `ctx := t.Context()` | + +```go +// before +func TestFetch(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := fetch(ctx) +} +// after +func TestFetch(t *testing.T) { + result := fetch(t.Context()) +} +``` + +### Go 1.24+ — `strings.Lines` / `bytes.Lines` + +| Before | After | +|---|---| +| `bufio.Scanner` line loop over a string/buffer | `for line := range strings.Lines(s)` | + +```go +// before +sc := bufio.NewScanner(strings.NewReader(s)) +for sc.Scan() { + process(sc.Text()) +} +// after +for line := range strings.Lines(s) { + process(strings.TrimSuffix(line, "\n")) +} +``` + +`strings.Lines` / `bytes.Lines` (Go 1.24) return line iterators — no Scanner setup, no default 64KB token-size limit. Note the yielded line **retains its trailing `\n`**, unlike `Scanner.Text()`; trim it if the old code relied on stripped lines. Only apply for in-memory strings/buffers, not streaming `io.Reader`s. + +### Go 1.24+ — `os.Root` (directory-scoped filesystem access) + +| Before | After | +|---|---| +| manual `filepath.Clean` + prefix check to block traversal | `root, _ := os.OpenRoot(dir); root.Open(name)` | + +```go +// before — hand-rolled path-traversal guard +p := filepath.Join(base, name) +if !strings.HasPrefix(filepath.Clean(p), filepath.Clean(base)+string(os.PathSeparator)) { + return errUnsafePath +} +f, err := os.Open(p) +// after — the OS enforces the boundary +root, err := os.OpenRoot(base) +if err != nil { return err } +defer root.Close() +f, err := root.Open(name) // symlinks/".." escaping base are rejected +``` + +`os.Root` (Go 1.24) confines all operations to a directory tree, rejecting `..` and symlink escapes at the syscall layer — far more robust than string prefix checks. **Security hardening:** flag as a strong suggestion wherever user-controlled paths are joined to a base directory. + +### Go 1.24+ — `omitzero` struct tag + +| Before | After | +|---|---| +| `json:"field,omitempty"` (for `time.Time`, `time.Duration`, structs, slices, maps) | `json:"field,omitzero"` | + +```go +// before +type Config struct { + Timeout time.Duration `json:"timeout,omitempty"` + Labels []string `json:"labels,omitempty"` +} +// after +type Config struct { + Timeout time.Duration `json:"timeout,omitzero"` + Labels []string `json:"labels,omitzero"` +} +``` + +Only for types where `omitempty` fails: `time.Time`, `time.Duration`, structs, slices, maps. Flag as suggestion, not auto-apply. + +### Go 1.24+ — `b.Loop()` in benchmarks + +| Before | After | +|---|---| +| `for i := 0; i < b.N; i++ { ... }` | `for b.Loop() { ... }` | + +```go +// before +func BenchmarkHash(b *testing.B) { + for i := 0; i < b.N; i++ { + hash(input) + } +} +// after +func BenchmarkHash(b *testing.B) { + for b.Loop() { + hash(input) + } +} +``` + +### Go 1.25+ — `sync.WaitGroup.Go` + +| Before | After | +|---|---| +| `wg.Add(1); go func() { defer wg.Done(); fn() }()` | `wg.Go(fn)` | + +```go +// before +var wg sync.WaitGroup +for _, item := range items { + wg.Add(1) + go func(item Item) { + defer wg.Done() + process(item) + }(item) +} +wg.Wait() +// after +var wg sync.WaitGroup +for _, item := range items { + wg.Go(func() { process(item) }) +} +wg.Wait() +``` + +### Go 1.26+ — `new` with expressions + +| Before | After | +|---|---| +| `v := val; &v` | `new(val)` | +| Helper `func ptr[T any](v T) *T { return &v }` | `new(val)` directly | + +```go +// before +timeout := 30 +debug := true +cfg := Config{ + Timeout: &timeout, + Debug: &debug, +} +// after +cfg := Config{ + Timeout: new(30), + Debug: new(true), +} +``` + +```go +// before +func ptr[T any](v T) *T { return &v } +cfg := Config{Count: ptr(10)} + +// after +cfg := Config{Count: new(10)} +``` + +### Go 1.26+ — `errors.AsType` + +| Before | After | +|---|---| +| `var t *T; errors.As(err, &t)` | `t, ok := errors.AsType[*T](err)` | + +```go +// before +var pathErr *os.PathError +if errors.As(err, &pathErr) { + log.Println(pathErr.Path) +} +// after +if pathErr, ok := errors.AsType[*os.PathError](err); ok { + log.Println(pathErr.Path) +} +``` + +### Go 1.27+ — Embedded field literals (embedlit) + +| Before | After | +|---|---| +| `T{U: U{x: 1}}` (redundant embedded-type specifier) | `T{x: 1}` | + +```go +type Base struct { + ID int + Name string +} +type User struct { + Base + Age int +} + +// before +u := User{ + Base: Base{ID: 1, Name: "alice"}, + Age: 30, +} +// after +u := User{ + ID: 1, + Name: "alice", + Age: 30, +} +``` + +The `embedlit` modernizer (gopls v0.22.0, `EmbedLitAnalyzer`) strips redundant embedded-struct field-type specifiers from composite literals. Go 1.27 lets you initialize promoted fields directly without the nested literal. Only apply when the promoted field names don't collide with the outer struct's own fields. + +## Operation Phases + +### Phase 1: Detect +Read `go.mod` to extract the Go version (`go 1.xx` line). If no `go.mod` is found, default to `go 1.21`. + +### Phase 2: Gather files +Find all `.go` files in the target scope (project root, or user-specified file/directory). Exclude `vendor/`, `.git/`, and `testdata/` directories. + +### Phase 3: Apply transformations +For each `.go` file, apply all transformations for versions ≤ the detected Go version. Process files sequentially. For each file: +1. Read the file content. +2. Identify applicable transformations by scanning for the "Before" patterns. +3. Apply each transformation using the Edit tool. +4. Run `goimports -w` (or `gofmt -w`) on the file after all edits. + +Never apply a transformation that requires a version higher than the project's Go version. + +### Phase 4: Report +Print a summary table showing: +- **File**: path relative to project root +- **Transformations applied**: list of transformation names per file +- **Total files modified** and **total transformations applied** +- **Skipped transformations** (available but not applicable due to version constraints) and their required Go version + +## Example Summary Output + +``` +## Modernization Summary + +| File | Transformations | +|---|---| +| main.go | any, strings.Cut, min/max (2 occurrences) | +| pkg/handler.go | range over int (3), slices.Contains, t.Context() | +| pkg/util.go | new(expr) (1), errors.AsType → errors.Is | + +**3 files modified, 10 transformations applied** + +Skipped (requires higher Go version): +- new(expr): requires go 1.26 (project is go 1.24) +- WaitGroup.Go: requires go 1.25 (project is go 1.24) +``` + +## Safety Rules + +- Never apply transformations that change semantics in edge cases without the user's awareness. +- Do not apply `omitzero` blindly—it changes JSON serialization behavior; flag it as a suggestion instead. +- Treat semantic migrations as suggestions, not auto-applies: `math/rand/v2` (changes the random stream), `iter.Seq` iterators (reshapes the API), `os.Root` (behavior/error-path change). Point them out and let the user opt in. +- When replacing a `bufio.Scanner` loop with `strings.Lines`/`bytes.Lines`, remember the yielded line keeps its trailing `\n`; add a `TrimSuffix` if the old code relied on `Scanner.Text()` semantics. +- Do not apply `strings.SplitSeq` or `bytes.SplitSeq` when the loop body references the index or the full slice elsewhere. +- Do not apply `strings.Builder` if the concatenation happens outside a loop (single `+=` is fine). +- When a transformation requires a new import, ensure the import is added to the file. +- After all edits, run `goimports -w` on each modified file to clean up imports. +- If `goimports` is not available, fall back to `gofmt -w`. +- If the project has no `go.mod`, ask the user for the target Go version before proceeding. + +## Automated tooling (Go 1.27+) + +Many of these transformations are now shipped as official modernizers in `gopls`/`go fix`. On Go 1.27+ the whole set can be applied across a module with: + +```sh +go fix ./... +``` + +gopls v0.22.0 added four notable passes covered above: + +| Modernizer | Min Go | Transformation | +|---|---|---| +| `unsafefuncs` | 1.17 | `uintptr` pointer math → `unsafe.Add` / `unsafe.Slice` | +| `atomictypes` | 1.19 | primitive `atomic.*` funcs → typed `atomic.Int32`/`Pointer[T]` wrappers | +| `slicesbackward` | 1.23 | descending-index loops → `slices.Backward` iterator | +| `embedlit` | 1.27 | redundant embedded-field literals → promoted-field init | + +Other modernizers in the same suite that this catalog covers: `minmax`, `efaceany`, `fmtappendf`, `stringscut`, `stringsseq`, `sortslice`/`slicescontains`, `mapsloop`, `stditerators`, `forvar`, `rangeint`, `testingcontext`, `bloop`, `waitgroup`, `newexpr`, `errorsastype`, `appendclipped` (→ `slices.Concat`), and `plusbuild` (→ drop obsolete `// +build`). + +To disable an over-eager pass (e.g. `slicesbackward` rewriting loops that mutate the slice), scope the run with `-fixes` or exclude that analyzer in your editor's gopls settings. diff --git a/pigo/internal/builtinskills/skills/note-it/SKILL.md b/pigo/internal/builtinskills/skills/note-it/SKILL.md new file mode 100644 index 0000000..da6a5e4 --- /dev/null +++ b/pigo/internal/builtinskills/skills/note-it/SKILL.md @@ -0,0 +1,188 @@ +--- +name: note-it +description: "Capture implementation notes after code implementation and review/fix. Records design decisions, deviations, tradeoffs, and open questions to docs/issue#XXXX.html. Triggers on: /note-it, 记录笔记, implementation notes." +user-invocable: true +--- + +# Implementation Notes + +After completing implementation and review/fix for an Issue, capture a running implementation notes file that documents how the implementation diverges from or interprets the spec. + +## Triggers + +Use when: +- After `/goal` implementation and `/review-it` are both complete +- User says "记录笔记", "implementation notes", "note-it", "/note-it" +- Before `/ship-it` (as a final checkpoint) +- Any time the user wants to capture design rationale + +## The Job + +1. Determine the Issue number from context (branch name, `/goal` target, or user input) +2. Review the implementation against the Issue spec / PRD +3. Generate an HTML notes file at `docs/issue#XXXX.html` +4. Present a summary to the user + +## Notes Structure + +The HTML file must cover these four categories. If a category has nothing to report, write "None" with a brief explanation. + +### 1. Design Decisions +Choices made where the spec was ambiguous or silent: +- What was the ambiguity? +- What choice did you make? +- What was the rationale? + +### 2. Deviations +Places where you intentionally departed from the spec: +- What did the spec say? +- What did you implement instead? +- Why was the deviation necessary or better? + +### 3. Tradeoffs +Alternatives you considered and why you picked what you did: +- What were the viable alternatives? +- What were the pros/cons of each? +- Why did the chosen approach win? + +### 4. Open Questions +Anything you'd want confirmed or revised: +- What assumption are you unsure about? +- What should the user verify? +- What might need follow-up? + +## Output + +- **Format:** HTML +- **Location:** `docs/` +- **Filename:** `issue#XXXX.html` (where XXXX is the zero-padded Issue number, e.g., `issue#0042.html`) + +## HTML Template + +Use this exact HTML structure: + +```html + + + + + + Implementation Notes — Issue #XXXX + + + +
+

Implementation Notes

+

Issue #{{ISSUE_NUMBER}} — {{ISSUE_TITLE}} — {{DATE}}

+ +

Design Decisions

+ + +

Deviations

+ + +

Tradeoffs

+ + +

Open Questions

+ + + +
+ + +``` + +## Example Item + +```html +
+

Decision Used interface-based polymorphism instead of switch

+

Ambiguity: The spec said "handle different types" without specifying how.

+

Choice: Defined a Handler interface with per-type implementations.

+

Rationale: Adding new types requires no changes to existing code (Open/Closed Principle). A switch would grow unboundedly.

+
+``` + +## How to Determine the Issue Number + +1. If the user provides it directly (e.g., `/note-it #42`), use it +2. If on a branch named `feat/issue-42-*` or `fix/issue-42-*`, extract `42` +3. If the last `/goal` target was `#42`, use `42` +4. Otherwise, ask the user: "Which Issue number should I use for the notes file?" + +## Edge Cases + +| Scenario | Handling | +|----------|----------| +| No Issue number found | Ask the user to specify | +| `docs/` directory does not exist | Auto-create it | +| Notes file already exists for this Issue | Ask: "Update existing notes or overwrite?" — default to update (append new items) | +| No deviations or open questions | Write "None — implementation followed the spec as written." | +| Spec/PRD file not found | Note in Open Questions: "No PRD found at tasks/prd-*.md — verify against original requirements." | + +## Checklist + +Before saving: +- [ ] Issue number identified +- [ ] All four categories reviewed (even if some are "None") +- [ ] Design decisions explain rationale, not just what was done +- [ ] Deviations clearly contrast spec vs implementation +- [ ] Tradeoffs mention specific alternatives considered +- [ ] Open questions are actionable (user can answer yes/no or give direction) +- [ ] HTML is well-formed and renders correctly \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/prd-to-spec/SKILL.md b/pigo/internal/builtinskills/skills/prd-to-spec/SKILL.md new file mode 100644 index 0000000..4d83ce5 --- /dev/null +++ b/pigo/internal/builtinskills/skills/prd-to-spec/SKILL.md @@ -0,0 +1,409 @@ +--- +name: prd-to-spec +description: "Transform a PRD into a technical SPEC document — architecture, API design, data model, error handling, and implementation contracts. Triggers on: prd-to-spec, prd to spec, prd转spec, 需求转设计, 需求转规格, generate spec from prd, design from prd, 技术方案, 设计方案." +user-invocable: true +--- + +# prd-to-spec — PRD to Technical Specification + +Transform a Product Requirements Document (PRD) into a detailed technical SPEC that an engineer or AI agent can implement against. The PRD says *what* to build; the SPEC says *how* to build it. + +--- + +## When to Use + +- A `/prd` has been generated and you need to bridge the gap to implementation +- You want architecture decisions documented before coding starts +- Multiple developers/agents will implement the feature and need a shared contract +- You need to validate technical feasibility before committing to a PRD +- You want to catch design issues early — before code is written + +--- + +## The Job + +1. **Locate PRD** — find or receive the PRD document +2. **Analyze context (optional)** — if a codebase exists, scan it to understand current architecture, patterns, and constraints +3. **Ask clarifying questions** — resolve technical ambiguities (max 3-5 questions) +4. **Generate SPEC** — produce a structured technical specification +5. **Review** — present to user for feedback and iteration +6. **Save** — write final SPEC to agreed location + +--- + +## Step 1: Locate PRD + +Find the input PRD in one of these ways: + +``` +Provide the PRD to convert: + +A. File path (e.g., tasks/prd-priority-system.md) +B. GitHub Issue URL +C. Paste PRD content directly +D. Auto-detect: scan tasks/ directory for recent PRDs +``` + +If auto-detecting, list available PRDs and let the user choose: + +``` +Found PRDs in tasks/: + 1. tasks/prd-priority-system.md (2024-03-15) + 2. tasks/prd-user-auth.md (2024-03-10) + +Which PRD should I convert? [1/2] +``` + +--- + +## Step 2: Analyze Context (Optional) + +**Skip this step if no codebase exists yet** (greenfield project). In that case, the SPEC will propose architecture from scratch based on the PRD requirements and clarifying questions. + +If a codebase exists, scan it to understand: + +- **Existing architecture** — how the current system is structured +- **Tech stack** — languages, frameworks, libraries already in use +- **Patterns** — naming conventions, file organization, error handling approach +- **Database** — current schema, migration tool, ORM +- **API style** — REST/GraphQL/gRPC, authentication method, response format +- **Testing** — test framework, coverage patterns, test utilities + +This ensures the SPEC aligns with the existing system rather than proposing incompatible solutions. + +--- + +## Step 3: Clarifying Questions + +Ask only when the PRD leaves technical decisions ambiguous. Focus on: + +- **Architecture choices** — where does this feature live? New service or extend existing? +- **Data storage** — new table? Extend existing? Cache strategy? +- **API design** — new endpoints? Extend existing? Breaking changes? +- **Dependencies** — any new libraries needed? Version constraints? +- **Performance** — expected load? Latency requirements? Batch size limits? + +Format: +``` +Technical questions before I generate the SPEC: + +1. Where should the priority logic live? + A. Extend existing TaskService + B. New PriorityService + C. Inline in controller + D. Let me decide based on the codebase + +2. Database migration approach? + A. Add column to existing tasks table + B. New priority table with FK + C. JSON field on tasks + D. Let me decide based on current schema + +3. API versioning concern? + A. Add to existing v1 endpoints + B. New v2 endpoints + C. No versioning needed +``` + +If user selects "let me decide" options, make the best choice based on codebase analysis and document the rationale in the SPEC. + +--- + +## Step 4: SPEC Document Structure + +```markdown +# SPEC: [Feature Name] + +> Technical specification derived from: [PRD filename/link] +> Generated: [date] | Target branch: [branch] | Commit: [short-hash] + +## 1. Summary + +### 1.1 What This SPEC Covers +[One paragraph: what feature this specifies and the scope of implementation] + +### 1.2 PRD Reference +- Source: [path or URL to PRD] +- User Stories covered: [US-001, US-002, ...] +- Functional Requirements covered: [FR-1, FR-2, ...] + +### 1.3 Design Decisions Summary +| Decision | Choice | Rationale | +|----------|--------|-----------| +| ... | ... | ... | + +--- + +## 2. Architecture + +### 2.1 System Context +[Where this feature fits in the overall system — diagram or description] + +### 2.2 Component Design +[New components/modules introduced, their responsibilities, and boundaries] + +### 2.3 Module Interactions +[How new components interact with existing ones — sequence or data flow] + +### 2.4 File Structure +[New files to create and existing files to modify] + +``` +src/ +├── services/ +│ └── priority.service.ts [NEW] +├── controllers/ +│ └── task.controller.ts [MODIFY: add priority endpoints] +├── models/ +│ └── priority.model.ts [NEW] +└── migrations/ + └── 20240315_add_priority.ts [NEW] +``` + +--- + +## 3. Data Model + +### 3.1 Schema Changes +[New tables, columns, indexes — with SQL or ORM notation] + +### 3.2 Entity Definitions +[TypeScript interfaces / Go structs / Python dataclasses for new entities] + +### 3.3 Relationships +[How new entities relate to existing ones — FK, embedded, reference] + +### 3.4 Migration Plan +[Migration steps, backward compatibility, rollback strategy] + +--- + +## 4. API Design + +### 4.1 Endpoints + +| Method | Path | Description | Auth | Request | Response | +|--------|------|-------------|------|---------|----------| +| ... | ... | ... | ... | ... | ... | + +### 4.2 Request/Response Schemas +[Detailed shapes with field types, validation rules, and examples] + +### 4.3 Error Responses +[Error codes, messages, and HTTP status codes for each failure mode] + +### 4.4 Breaking Changes +[Any backward-incompatible changes and migration path for consumers] + +--- + +## 5. Business Logic + +### 5.1 Core Algorithms +[Step-by-step logic for key operations — pseudocode or structured description] + +### 5.2 Validation Rules +[Input validation, business rule validation, with specific constraints] + +### 5.3 State Machine +[If applicable: states, transitions, guards, and side effects] + +### 5.4 Edge Cases +[Known edge cases and how they should be handled] + +--- + +## 6. Error Handling + +### 6.1 Error Taxonomy +| Error Code | HTTP Status | Condition | User Message | +|------------|-------------|-----------|--------------| +| ... | ... | ... | ... | + +### 6.2 Retry Strategy +[Which operations are retryable, backoff policy, max attempts] + +### 6.3 Failure Modes +[What happens when dependencies fail — graceful degradation plan] + +--- + +## 7. Security + +### 7.1 Authentication & Authorization +[Who can access what, permission model, role checks] + +### 7.2 Input Validation +[Sanitization rules, injection prevention, size limits] + +### 7.3 Data Protection +[Sensitive fields, encryption at rest/transit, audit logging] + +--- + +## 8. Performance + +### 8.1 Expected Load +[Estimated QPS, data volume, growth projection] + +### 8.2 Optimization Strategy +[Caching, pagination, lazy loading, batch processing] + +### 8.3 Database Considerations +[Index strategy, query patterns, N+1 prevention] + +--- + +## 9. Testing Strategy + +### 9.1 Unit Tests +[What to test, test boundaries, mock strategy] + +### 9.2 Integration Tests +[API tests, database tests, service interaction tests] + +### 9.3 Edge Case Tests +[Specific scenarios to cover based on Section 5.4] + +### 9.4 Acceptance Criteria Mapping +| US/FR | Test | Type | Description | +|-------|------|------|-------------| +| US-001 | ... | unit | ... | +| FR-2 | ... | integration | ... | + +--- + +## 10. Implementation Plan + +### 10.1 Phases +[Order of implementation — what to build first, dependencies between steps] + +### 10.2 Issue Mapping +[Map SPEC sections to PRD Issues for implementation tracking] + +| Issue | SPEC Sections | Priority | Depends On | +|-------|--------------|----------|------------| +| #1 | 3.1, 3.4 | high | — | +| #2 | 4.1, 4.2, 5.1 | high | #1 | +| ... | ... | ... | ... | + +### 10.3 Incremental Delivery +[How to ship incrementally — feature flags, dark launches, gradual rollout] + +--- + +## 11. Open Questions & Risks + +### 11.1 Unresolved Questions +- [Questions that need product/engineering input before implementation] + +### 11.2 Technical Risks +| Risk | Impact | Mitigation | +|------|--------|-----------| +| ... | ... | ... | + +### 11.3 Assumptions +- [Technical assumptions made during SPEC creation — validate before implementing] +``` + +--- + +## Step 5: Review & Iteration + +After generating the SPEC, present it and ask: + +``` +SPEC generated from PRD. Please review: + +- Are the architecture choices appropriate? +- Are there missing edge cases or error scenarios? +- Is the API design consistent with existing patterns? +- Should any section have more/less detail? + +Reply OK to save, or provide feedback for iteration. +``` + +--- + +## Step 6: Save + +Ask user for save location: + +``` +Where should I save the SPEC? + +A. tasks/spec-[feature-name].md (alongside PRD, recommended) +B. docs/spec-[feature-name].md +C. Custom path: [specify] +``` + +--- + +## Mapping Strategy: PRD → SPEC + +How PRD elements translate to SPEC sections: + +| PRD Section | SPEC Section(s) | Transformation | +|-------------|-----------------|----------------| +| User Stories | 5. Business Logic, 9.4 Acceptance Mapping | Stories → algorithms + test cases | +| Functional Requirements | 4. API Design, 5. Business Logic | FRs → endpoints + logic | +| Acceptance Criteria | 9. Testing Strategy | Criteria → specific test scenarios | +| Non-Goals | 11.1 Open Questions | Clarify what's explicitly excluded | +| Technical Considerations | 2. Architecture, 8. Performance | Constraints → design decisions | +| Success Metrics | 8.1 Expected Load, 10.3 Delivery | Metrics → monitoring + rollout plan | + +--- + +## Quality Criteria + +A good SPEC should pass these checks: + +- [ ] Every PRD User Story has corresponding SPEC sections +- [ ] Every Functional Requirement maps to an API endpoint or business logic rule +- [ ] Every Acceptance Criterion maps to at least one test case +- [ ] Architecture choices are justified with rationale +- [ ] API schemas are specific enough to generate client code +- [ ] Error handling covers all identified failure modes +- [ ] Implementation order respects dependencies +- [ ] No "TBD" or "TODO" items — resolve or move to Open Questions + +--- + +## Edge Cases & Fallback + +| Scenario | Handling | +|----------|----------| +| PRD is vague or incomplete | Generate SPEC with best-effort choices, mark assumptions in Section 11.3 | +| PRD conflicts with existing code | Flag conflicts explicitly, propose resolution in Section 11.1 | +| Feature is too large for one SPEC | Split into multiple SPECs (one per service boundary), link them | +| No existing codebase (greenfield) | Skip Step 2, propose architecture from scratch based on PRD + clarifying questions | +| PRD has no User Stories (just bullet points) | Infer structure, map bullets to SPEC sections, note in Summary | +| User wants SPEC without reading codebase | Skip Step 2, note that assumptions about existing code are unverified | +| Multiple PRDs need one SPEC | Merge PRD inputs, deduplicate requirements, note source for each | + +--- + +## Anti-Patterns to Avoid + +- **Don't restate the PRD.** The SPEC adds technical depth, not a copy of requirements in different words. +- **Don't over-specify trivial operations.** CRUD with no special logic doesn't need a full algorithm section. +- **Don't pick technologies without context.** Always check what the project already uses before suggesting new tools. +- **Don't design in isolation.** The SPEC must fit the existing system — same patterns, same conventions, same style. +- **Don't leave decisions implicit.** If you made a choice (e.g., "add column to existing table"), state it and say why. +- **Don't write implementation code.** The SPEC describes contracts and behavior, not code. Pseudocode is acceptable for complex algorithms. + +--- + +## Relationship to Other Skills + +``` +/prd → /prd-to-spec → /goal → /review-it → /ship-it + │ │ │ + │ Requirements │ Technical │ Implementation + │ (what) │ (how) │ (code) +``` + +- **/prd** produces the PRD (input to this skill) +- **/prd-to-spec** produces the SPEC (this skill) +- **/goal** implements Issues with SPEC as the technical reference +- **/code-to-spec** reverse-engineers SPEC from existing code (complementary — forward vs. reverse) diff --git a/pigo/internal/builtinskills/skills/prd/LICENSE b/pigo/internal/builtinskills/skills/prd/LICENSE new file mode 100644 index 0000000..14fac91 --- /dev/null +++ b/pigo/internal/builtinskills/skills/prd/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/pigo/internal/builtinskills/skills/prd/README.md b/pigo/internal/builtinskills/skills/prd/README.md new file mode 100644 index 0000000..6738757 --- /dev/null +++ b/pigo/internal/builtinskills/skills/prd/README.md @@ -0,0 +1,43 @@ +# PRD Generator Skill + +Generate structured Product Requirements Documents (PRD) for new features. Focused solely on producing a clear, implementable PRD — Issue decomposition and technical design are handled by separate skills. + +## Features + +- Asks 3-5 clarifying questions with lettered options for quick iteration +- Generates a well-structured PRD with user stories, numbered functional requirements, non-goals, success metrics, and more +- Enforces verifiable acceptance criteria (observable / testable / verifiable) +- Supports user review and adjustment before saving +- Saves output to `tasks/prd-[feature-name].md` +- Bilingual (Chinese & English) edge case handling + +## Workflow + +The PRD skill is the first step in a three-stage pipeline: + +| Stage | Skill | Purpose | +|-------|-------|---------| +| 1. Requirements | `/prd` (this skill) | Define *what* to build | +| 2. Technical design (optional) | `/prd-to-spec` | Define *how* to build it | +| 3. Decomposition | `/to-issues` | Break into implementable tickets (GitHub / Local / Baidu iCafe) | + +After a PRD is confirmed, run `/prd-to-spec` for complex features, then `/to-issues` — or go straight to `/to-issues`. + +## Usage + +Trigger with prompts like: + +- "create a prd for..." +- "write prd for..." +- "写PRD" +- "需求文档" +- "需求分析" + +## Files + +- `SKILL.md` — Skill definition and instructions +- `test-prompts.json` — Test prompts for validation + +## Attribution + +This skill is adapted from [ralph/skills/prd](https://github.com/snarktank/ralph/tree/main/skills/prd). diff --git a/pigo/internal/builtinskills/skills/prd/SKILL.md b/pigo/internal/builtinskills/skills/prd/SKILL.md new file mode 100644 index 0000000..c682da4 --- /dev/null +++ b/pigo/internal/builtinskills/skills/prd/SKILL.md @@ -0,0 +1,278 @@ +--- +name: prd +description: "Generate a Product Requirements Document (PRD) for a new feature. Use when planning a feature, starting a new project, or when asked to create a PRD. After PRD is confirmed, use /prd-to-spec (optional) for technical design, then /to-issues to create implementable tickets. Triggers on: create a prd, write prd for, plan this feature, requirements for, spec out, 写PRD, 需求文档, 需求分析, 规格说明." +user-invocable: true +--- + +# PRD Generator + +Create detailed Product Requirements Documents that are clear, actionable, and suitable for implementation. After PRD is confirmed, use `/to-issues` to decompose it into Issues, and optionally `/prd-to-spec` for technical design before that. + +--- + +## The Job + +1. Receive a feature description from the user +2. Ask clarifying questions to cover key ambiguities — scale the count to complexity, not a fixed number (see Step 1) +3. Generate a structured PRD based on answers +4. **Present PRD to user for review** — ask "Please review the PRD. Let me know if any adjustments are needed, or reply OK to confirm." +5. Apply any adjustments, then save to `tasks/prd-[feature-name].md` +6. **Suggest next steps** (see Step 3) + +**Important:** Do NOT start implementing. Just create the PRD. + +--- + +## Step 1: Clarifying Questions + +Ask only critical questions where the initial prompt is ambiguous. **Scale the number of questions to the feature's complexity — the goal is covering key ambiguities, not hitting a fixed count:** + +- **Simple, well-scoped feature:** 2-3 questions +- **Typical feature:** 3-5 questions +- **Complex feature** (multiple user roles, cross-system integration, significant ambiguity): 6-8 questions + +If a dimension is already unambiguous from the user's input, skip it — don't ask filler questions just to reach a number. Focus on: + +- **Problem/Goal:** What problem does this solve? +- **Core Functionality:** What are the key actions? +- **Scope/Boundaries:** What should it NOT do? +- **Success Criteria:** How do we know it's done? + +### Format Questions Like This: + +``` +1. What is the primary goal of this feature? + A. Improve user onboarding experience + B. Increase user retention + C. Reduce support burden + D. Other: [please specify] + +2. Who is the target user? + A. New users only + B. Existing users only + C. All users + D. Admin users only + +3. What is the scope? + A. Minimal viable version + B. Full-featured implementation + C. Just the backend/API + D. Just the UI +``` + +This lets users respond with "1A, 2C, 3B" for quick iteration. Remember to indent the options. + +--- + +## Edge Cases & Fallback + +| Scenario | Handling | +|----------|----------| +| User skips clarifying questions (e.g., replies "whatever", "just write it") | Fill with reasonable defaults, mark with `[Assumption]` in PRD, prompt user to confirm during review | +| User input is too vague (e.g., "add a feature") | Ask once for specifics; if still vague, infer from project context and mark assumptions | +| `tasks/` directory does not exist | Auto-create `tasks/` directory | +| feature-name is hard to extract from input | Ask the user directly: "Suggested PRD filename is prd-XXX.md, please confirm or modify" | +| User requests PRD changes after review | Apply changes and re-save without re-running the clarification flow | +| PRD content exceeds 500 lines | Suggest the user consider splitting into multiple sub-feature PRDs | +| User declines to proceed | Just save the PRD, user can run `/to-issues` later | +| Issue creation needed later | Suggest running `/to-issues` with the saved PRD file | + +--- + +## Step 2: PRD Structure + +Generate the PRD with these sections: + +### 1. Introduction/Overview +Brief description of the feature and the problem it solves. Use plain language — avoid jargon or explain it. Assume the reader may be a junior developer or AI agent. + +### 2. Goals +Specific, measurable objectives (bullet list). + +### 3. User Stories +Each story needs: +- **Title:** Short descriptive name +- **Description:** "As a [user], I want [feature] so that [benefit]" +- **Acceptance Criteria:** Verifiable checklist of what "done" means + +**Numbering rule:** US-001, US-002, US-003... (three digits, starting from 001). Each US should be independently implementable and small enough to complete within one focused agent session. + +**Acceptance criteria self-check template:** Each criterion must satisfy at least one of the following, otherwise it is considered "vague" and must be rewritten: +- Observable: describes a specific UI state or API response (e.g., "button shows confirmation dialog") +- Testable: has clear input/output pairs (e.g., "entering an empty email shows a red warning") +- Verifiable: can be checked by tools (e.g., "Typecheck/lint passes") +- ❌ Bad example: "works correctly", "good user experience", "excellent performance" → these are unverifiable + +**Format:** +```markdown +### US-001: [Title] +**Description:** As a [user], I want [feature] so that [benefit]. + +**Acceptance Criteria:** +- [ ] Specific verifiable criterion +- [ ] Another criterion +- [ ] Typecheck/lint passes +- [ ] **[UI stories only]** Verify in a browser (e.g., via the `run` skill) +``` + +**Important:** +- Acceptance criteria must be verifiable, not vague. "Works correctly" is bad. "Button shows confirmation dialog before deleting" is good. +- **For any story with UI changes:** Always include "Verify in a browser" as acceptance criteria (e.g., via the `run` skill). This ensures visual verification of frontend work. + +### 4. Functional Requirements +Numbered list of specific functionalities: +- "FR-1: The system must allow users to..." +- "FR-2: When a user clicks X, the system must..." + +**FR specification:** Each FR starts with `FR-N:` (N increments from 1), uses "system must / system shall" phrasing, and describes **one** specific behavior. Avoid combining multiple "and"-linked behaviors in a single FR. + +### 5. Non-Goals (Out of Scope) +What this feature will NOT include. Critical for managing scope. + +### 6. Design Considerations (Optional) +- UI/UX requirements +- Link to mockups if available +- Relevant existing components to reuse + +### 7. Technical Considerations (Optional) +- Known constraints or dependencies +- Integration points with existing systems +- Performance requirements + +### 8. Success Metrics +How will success be measured? +- "Reduce time to complete X by 50%" +- "Increase conversion rate by 10%" + +### 9. Open Questions +Remaining questions or areas needing clarification. + +--- + +## Output + +- **Format:** Markdown (`.md`) +- **Location:** `tasks/` +- **Filename:** `prd-[feature-name].md` (kebab-case) + +--- + +## Step 3: Next Steps + +After the PRD is saved, suggest the user: + +``` +✅ PRD saved to tasks/prd-[feature-name].md + +Next steps: + /prd-to-spec → Generate technical SPEC (optional — for complex features) + /to-issues → Decompose into Issues and create tickets + +Or go straight to implementation: + /to-issues → Create Issues, then /goal to implement +``` + +If the user wants to proceed, invoke the corresponding skill. + +--- + +## Example PRD + +```markdown +# PRD: Task Priority System + +## Introduction + +Add priority levels to tasks so users can focus on what matters most. Tasks can be marked as high, medium, or low priority, with visual indicators and filtering to help users manage their workload effectively. + +## Goals + +- Allow assigning priority (high/medium/low) to any task +- Provide clear visual differentiation between priority levels +- Enable filtering and sorting by priority +- Default new tasks to medium priority + +## User Stories + +### US-001: Add priority field to database +**Description:** As a developer, I need to store task priority so it persists across sessions. + +**Acceptance Criteria:** +- [ ] Add priority column to tasks table: 'high' | 'medium' | 'low' (default 'medium') +- [ ] Generate and run migration successfully +- [ ] Typecheck passes + +### US-002: Display priority indicator on task cards +**Description:** As a user, I want to see task priority at a glance so I know what needs attention first. + +**Acceptance Criteria:** +- [ ] Each task card shows colored priority badge (red=high, yellow=medium, gray=low) +- [ ] Priority visible without hovering or clicking +- [ ] Typecheck passes +- [ ] Verify in a browser (e.g., via the `run` skill) + +### US-003: Add priority selector to task edit +**Description:** As a user, I want to change a task's priority when editing it. + +**Acceptance Criteria:** +- [ ] Priority dropdown in task edit modal +- [ ] Shows current priority as selected +- [ ] Saves immediately on selection change +- [ ] Typecheck passes +- [ ] Verify in a browser (e.g., via the `run` skill) + +### US-004: Filter tasks by priority +**Description:** As a user, I want to filter the task list to see only high-priority items when I'm focused. + +**Acceptance Criteria:** +- [ ] Filter dropdown with options: All | High | Medium | Low +- [ ] Filter persists in URL params +- [ ] Empty state message when no tasks match filter +- [ ] Typecheck passes +- [ ] Verify in a browser (e.g., via the `run` skill) + +## Functional Requirements + +- FR-1: Add `priority` field to tasks table ('high' | 'medium' | 'low', default 'medium') +- FR-2: Display colored priority badge on each task card +- FR-3: Include priority selector in task edit modal +- FR-4: Add priority filter dropdown to task list header +- FR-5: Sort by priority within each status column (high to medium to low) + +## Non-Goals + +- No priority-based notifications or reminders +- No automatic priority assignment based on due date +- No priority inheritance for subtasks + +## Technical Considerations + +- Reuse existing badge component with color variants +- Filter state managed via URL search params +- Priority stored in database, not computed + +## Success Metrics + +- Users can change priority in under 2 clicks +- High-priority tasks immediately visible at top of lists +- No regression in task list performance + +## Open Questions + +- Should priority affect task ordering within a column? +- Should we add keyboard shortcuts for priority changes? +``` + +--- + +## Checklist + +Before saving the PRD: + +- [ ] Asked clarifying questions with lettered options +- [ ] Incorporated user's answers +- [ ] User stories are small and specific +- [ ] Functional requirements are numbered and unambiguous +- [ ] Non-goals section defines clear boundaries +- [ ] Saved to `tasks/prd-[feature-name].md` +- [ ] Suggested next steps: `/prd-to-spec` (optional) and `/to-issues` diff --git a/pigo/internal/builtinskills/skills/prd/test-prompts.json b/pigo/internal/builtinskills/skills/prd/test-prompts.json new file mode 100644 index 0000000..5413d13 --- /dev/null +++ b/pigo/internal/builtinskills/skills/prd/test-prompts.json @@ -0,0 +1,27 @@ +[ + { + "id": 1, + "prompt": "帮我写一个用户登录功能的PRD", + "expected": "应先询问3-5个澄清问题(登录方式、目标用户、安全要求等),用户回答后生成结构化PRD,包含用户故事、编号功能需求、非目标、验收标准等完整章节;结尾建议下一步运行 /prd-to-spec(可选)和 /to-issues" + }, + { + "id": 2, + "prompt": "我需要给电商App加一个购物车功能,支持多规格商品", + "expected": "应识别'多规格'的歧义(SKU选择 vs 自定义规格),在澄清问题中覆盖;PRD中需体现规格选择、数量修改、价格联动等核心逻辑,验收标准需可验证" + }, + { + "id": 3, + "prompt": "创建一个PRD,给我们的内部工具加批量导入功能", + "expected": "应区分内部工具场景(无普通C端用户),PRD应聚焦批量操作的错误处理、进度反馈、部分失败策略等企业场景关注点" + }, + { + "id": 4, + "prompt": "随便写个功能的PRD吧", + "expected": "输入过于模糊时应先追问一次具体功能;若用户仍不明确(如回复'whatever'),用合理默认值填充并在PRD中用[Assumption]标注,在review阶段提示用户确认" + }, + { + "id": 5, + "prompt": "给任务列表加一个优先级功能的PRD,包含前端展示", + "expected": "含UI改动的用户故事,验收标准必须包含'Verify in a browser (e.g., via the run skill)'一项;生成后展示PRD供用户review,确认后保存到 tasks/prd-*.md" + } +] diff --git a/pigo/internal/builtinskills/skills/refactor/SKILL.md b/pigo/internal/builtinskills/skills/refactor/SKILL.md new file mode 100644 index 0000000..2647f68 --- /dev/null +++ b/pigo/internal/builtinskills/skills/refactor/SKILL.md @@ -0,0 +1,674 @@ +--- +name: refactor +description: "Expert code refactoring based on Martin Fowler's catalog — improve maintainability without changing behavior. Covers code smells, composing methods, moving features, organizing data, simplifying conditionals, method calls, and generalization. Triggers on: refactor, 重构, clean up, improve code, code smell, extract method, rename, simplify." +user-invocable: true +--- + +# Refactor — Expert Code Restructuring + +Surgical code refactoring based on Martin Fowler's (2nd Edition) catalog. Improve structure, readability, and maintainability without changing external behavior. Gradual evolution, not revolution. + +--- + +## When to Use + +This skill activates when: +- Code is hard to understand or maintain +- Functions/classes have grown too large +- Code smells are detected +- Adding features is difficult due to poor structure +- User explicitly requests refactoring, cleanup, or improvement +- User says: refactor, 重构, clean up, improve code, code smell, extract method, rename, simplify + +--- + +## The Golden Rules + +These five rules are non-negotiable. Violating any of them turns refactoring into reckless editing. + +### 1. Behavior is Preserved + +Only *how* the code works changes, never *what* it does. If tests existed before, they must pass after. If the refactoring introduces a behavioral change, it's not refactoring — it's rewriting. + +### 2. Small Steps + +Each change should be the smallest possible transformation that compiles and passes tests. If a step breaks, you know exactly which change caused it. Refactoring is a series of tiny, safe transformations, not one big rewrite. + +### 3. Version Control is Your Friend + +Commit before starting. Commit after each successful step. This gives you infinite undo. Branch from a clean state so you can abandon the refactoring without consequences. + +### 4. Tests are Essential + +"Without tests, you're not refactoring — you're just editing." If tests don't exist for the target code, write characterization tests first. These tests capture the current behavior so you can detect regressions. + +### 5. One Thing at a Time + +Never mix refactoring with feature changes. Never refactor two unrelated things simultaneously. Each commit should contain exactly one refactoring operation. + +--- + +## When NOT to Refactor + +| Scenario | Action | +|----------|--------| +| Code works and won't change again | Leave it alone | +| Critical production path with no tests | Write characterization tests first | +| Under tight deadline pressure | Document the smell, refactor later | +| No clear purpose or benefit | Don't refactor for refactoring's sake | +| Code is fundamentally wrong | This is a rewrite, not a refactoring | + +--- + +## Code Smells Catalog + +Based on Fowler's taxonomy. Before refactoring, identify which smell is present. + +### Bloaters + +| Smell | Description | Primary Refactoring | +|-------|-------------|-------------------| +| **Long Method** | Method > 10-15 lines, doing multiple things | Extract Method, Replace Temp with Query | +| **Large Class** | Class with too many fields/methods (God Object) | Extract Class, Extract Subclass | +| **Primitive Obsession** | Using primitives instead of small objects | Replace Data Value with Object, Replace Type Code with Class | +| **Long Parameter List** | Method with > 3-4 parameters | Introduce Parameter Object, Preserve Whole Object | +| **Data Clumps** | Same group of data appearing together | Extract Class, Introduce Parameter Object | + +### Object-Orientation Abusers + +| Smell | Description | Primary Refactoring | +|-------|-------------|-------------------| +| **Switch Statements** | Repeated switch/if-else on type codes | Replace Conditional with Polymorphism, Replace Type Code with Subclasses | +| **Temporary Field** | Field only set in certain circumstances | Extract Class, Introduce Null Object | +| **Refused Bequest** | Subclass doesn't use inherited members | Replace Inheritance with Delegation, Push Down Method/Field | +| **Alternative Classes with Different Interfaces** | Classes doing similar things with different names | Rename Method, Move Method, Extract Superclass | + +### Change Preventers + +| Smell | Description | Primary Refactoring | +|-------|-------------|-------------------| +| **Divergent Change** | One class changed for different reasons | Extract Class | +| **Shotgun Surgery** | One change requires many small changes across classes | Move Method, Move Field, Inline Class | +| **Parallel Inheritance Hierarchies** | Adding a subclass to one hierarchy forces adding to another | Move Method, Move Field | + +### Dispensables + +| Smell | Description | Primary Refactoring | +|-------|-------------|-------------------| +| **Comments** | Comments explaining what code does (not why) | Extract Method, Rename Variable, Introduce Assertion | +| **Duplicate Code** | Same code structure in multiple places | Extract Method, Pull Up Method, Form Template Method | +| **Lazy Class** | Class doing too little to justify existence | Inline Class, Collapse Hierarchy | +| **Data Class** | Class with only fields and getters/setters | Move Method, Encapsulate Field, Encapsulate Collection | +| **Dead Code** | Unused code, imports, commented-out blocks | Delete it (git history has it) | +| **Speculative Generality** | Code built for "someday" that never came | Inline Class, Collapse Hierarchy, Remove Parameter | + +### Couplers + +| Smell | Description | Primary Refactoring | +|-------|-------------|-------------------| +| **Feature Envy** | Method uses another class's data more than its own | Move Method, Extract Method + Move Method | +| **Inappropriate Intimacy** | Classes know too much about each other's internals | Move Method, Move Field, Replace Delegation with Hidden Delegate | +| **Message Chains** | `a.getB().getC().getD().doSomething()` | Hide Delegate, Extract Method | +| **Middle Man** | Class delegates everything to another class | Remove Middle Man, Inline Method | +| **Incomplete Library Class** | Library missing methods you need | Introduce Foreign Method, Introduce Local Extension | + +--- + +## Refactoring Techniques Catalog + +Organized by category, from Fowler's catalog. Each technique includes its mechanical steps. + +### Composing Methods + +#### Extract Method +Turn a code fragment into a method whose name explains its purpose. + +**Mechanics:** +1. Create a new method named after what the fragment does (not how) +2. Copy the extracted code into the new method +3. Identify local variables: read-only become parameters, modified become return values +4. Pass parameters and handle return values +5. Replace the original fragment with a call to the new method +6. Test + +**Before:** +```java +void printOwing() { + printBanner(); + // Print details + System.out.println("name: " + _name); + System.out.println("amount: " + getOutstanding()); +} +``` + +**After:** +```java +void printOwing() { + printBanner(); + printDetails(getOutstanding()); +} + +void printDetails(double outstanding) { + System.out.println("name: " + _name); + System.out.println("amount: " + outstanding); +} +``` + +#### Inline Method +Replace a method call with its body when the method body is as clear as the name. + +**Mechanics:** +1. Check the method is not polymorphic (no subclasses override it) +2. Find all callers +3. Replace each call with the method body +4. Delete the method definition +5. Test + +#### Extract Variable +Put the result of an expression (or part of it) in a self-explanatory variable. + +**Before:** +```java +if (platform.toUpperCase().indexOf("MAC") > -1 && + browser.toUpperCase().indexOf("IE") > -1 && + wasInitialized() && resize > 0) { + // ... +} +``` + +**After:** +```java +final boolean isMacOs = platform.toUpperCase().indexOf("MAC") > -1; +final boolean isIEBrowser = browser.toUpperCase().indexOf("IE") > -1; +final boolean wasResized = resize > 0; +if (isMacOs && isIEBrowser && wasInitialized() && wasResized) { + // ... +} +``` + +#### Inline Temp +Replace a temp variable with its expression when the temp is only used once and the expression is clear. + +#### Replace Temp with Query +Extract the expression into a method. Temps that are computed once and reused are replaced with method calls. + +#### Split Temporary Variable +A temp assigned more than once (not loop/collecting) should be split into separate variables, one per responsibility. + +#### Remove Assignments to Parameters +Don't assign to parameters. Use a local variable instead. + +#### Replace Method with Method Object +When a long method uses many local variables that make Extract Method hard, turn the method into its own class, with locals as fields. + +#### Substitute Algorithm +Replace an algorithm with a clearer one. + +--- + +### Moving Features Between Objects + +#### Move Method +Move a method to the class where it's used most. + +**Mechanics:** +1. Check all features used by the method on its current class +2. Check for polymorphism (subclass/superclass methods) +3. Create the method on the target class, adapting as needed +4. Reference the target object from the source +5. Turn the source method into a delegating method, or remove it +6. Test + +#### Move Field +Move a field to the class where it's used most. + +#### Extract Class +When a class does the work of two, split it. Create a new class and move relevant fields and methods. + +#### Inline Class +When a class does almost nothing, absorb it into the class that uses it most. + +#### Hide Delegate +Create methods on the server to hide the delegate chain. `manager = person.getDepartment().getManager()` → `manager = person.getManager()`. + +#### Remove Middle Man +When a class is doing too much delegation, call the delegate directly. + +#### Introduce Foreign Method +When a server class needs an additional method but you can't modify it, create a method on the client with the server instance as the first argument. + +#### Introduce Local Extension +When you need multiple foreign methods, create an extension class (subclass or wrapper). + +--- + +### Organizing Data + +#### Self Encapsulate Field +Access fields through getters and setters, even within the owning class. + +#### Replace Data Value with Object +When a data item needs additional data or behavior, turn it into an object. + +**Before:** +```java +class Order { + private String customer; // Just a string +} +``` + +**After:** +```java +class Order { + private Customer customer; // Rich object with name, address, credit rating +} +``` + +#### Change Value to Reference +When you need to share one instance of an object across multiple places. + +#### Change Reference to Value +When a reference object is small, immutable, and you want value semantics. + +#### Replace Array with Object +When an array holds heterogeneous data (`String[] row = new String[3]` — name, score, wins), replace with an object. + +#### Duplicate Observed Data +Domain data lives in a GUI control but domain logic needs it. Copy the data into a domain object and set up an observer to keep the two in sync (Observer pattern). Separates presentation from domain so each can evolve independently. + +#### Change Unidirectional Association to Bidirectional +Two classes need each other's features but only one holds a reference. Add a back-pointer and make the modifiers on both ends keep the link consistent. Add the reference only when genuinely needed — bidirectional links raise coupling and risk inconsistency. + +#### Change Bidirectional Association to Unidirectional +A two-way link exists but one side no longer uses the other. Drop the unneeded direction. Reduces coupling, simplifies lifecycle management, and avoids "zombie" objects kept alive only by a stale back-pointer. + +#### Replace Magic Number with Symbolic Constant +Replace literal numbers/strings with named constants. + +#### Encapsulate Field +Make public fields private and provide accessors. + +#### Encapsulate Collection +Never return the raw collection. Return a read-only view and provide add/remove methods. + +#### Replace Type Code with Class +Replace a numeric/string type code with a class that has meaningful behavior. + +#### Replace Type Code with Subclasses +When type code affects behavior, use polymorphism instead of conditionals. + +#### Replace Type Code with State/Strategy +Similar to subclasses but uses composition when the type can change at runtime. + +#### Replace Subclass with Fields +When subclasses vary only in constant data, replace them with fields on a single class. + +--- + +### Simplifying Conditional Expressions + +#### Decompose Conditional +Extract the condition, then-part, and else-part into separate methods. + +**Before:** +```java +if (date.before(SUMMER_START) || date.after(SUMMER_END)) { + charge = quantity * _winterRate + _winterServiceCharge; +} else { + charge = quantity * _summerRate; +} +``` + +**After:** +```java +if (isSummer(date)) { + charge = summerCharge(quantity); +} else { + charge = winterCharge(quantity); +} +``` + +#### Consolidate Conditional Expression +Combine multiple conditionals that have the same result. + +#### Consolidate Duplicate Conditional Fragments +Move code that appears in every branch outside the conditional. + +#### Remove Control Flag +Replace control flags with break, continue, or return. + +#### Replace Nested Conditional with Guard Clauses +Use early returns for special cases instead of deep nesting. + +**Before (arrow code):** +```java +double getPayAmount() { + double result; + if (_isDead) { + result = deadAmount(); + } else { + if (_isSeparated) { + result = separatedAmount(); + } else { + if (_isRetired) { + result = retiredAmount(); + } else { + result = normalPayAmount(); + } + } + } + return result; +} +``` + +**After:** +```java +double getPayAmount() { + if (_isDead) return deadAmount(); + if (_isSeparated) return separatedAmount(); + if (_isRetired) return retiredAmount(); + return normalPayAmount(); +} +``` + +#### Replace Conditional with Polymorphism +When a conditional chooses different behavior based on the type of an object, use subclasses. + +#### Introduce Null Object +Replace null checks with a null object that provides default behavior. + +#### Introduce Assertion +State assumptions explicitly with assertions. + +--- + +### Making Method Calls Simpler + +#### Rename Method +The name should say what the method does. If you can't think of a good name, the method may have multiple responsibilities. + +#### Add Parameter / Remove Parameter +Add parameters when a method needs more info. Remove parameters when the method can get the info another way. + +#### Separate Query from Modifier +A method should either return a value OR change state, never both. + +#### Parameterize Method +Several methods doing similar things with different values → one method with a parameter. + +#### Replace Parameter with Explicit Methods +The inverse: when a parameter essentially selects different behavior, create separate methods. + +#### Preserve Whole Object +Pass the whole object instead of pulling individual fields from it. + +#### Replace Parameter with Method +*(refactoring.guru: Replace Parameter with Method Call)* When a parameter can be computed from data the object already has, remove the parameter and let the method call the query itself. + +#### Introduce Parameter Object +Group parameters that naturally go together into an object. + +#### Remove Setting Method +Make a field immutable by removing its setter and setting it in the constructor. + +#### Hide Method +Make methods private when they're not used outside the class. + +#### Replace Constructor with Factory Method +When you need more flexibility than a simple constructor call. + +#### Replace Error Code with Exception +Throw an exception instead of returning an error code. + +#### Replace Exception with Test +Check the condition first instead of catching an exception. + +--- + +### Dealing with Generalization + +#### Pull Up Field/Method/Constructor Body +Move identical fields/methods/constructor code from subclasses to superclass. + +#### Push Down Method/Field +Move behavior from superclass to only the subclasses that use it. + +#### Extract Subclass +Create a subclass for a subset of features used in some instances. + +#### Extract Superclass +Create a superclass for shared features of similar classes. + +#### Extract Interface +Create an interface from a subset of a class's public methods. + +#### Collapse Hierarchy +Merge a superclass and subclass when they're not different enough. + +#### Form Template Method +Generalize an algorithm in the superclass, letting subclasses fill in the specifics. + +#### Replace Inheritance with Delegation +When a subclass only uses part of the superclass, use composition instead. + +#### Replace Delegation with Inheritance +When a delegating class needs access to all of the delegate's behavior. + +--- + +## The Refactoring Process + +### Phase 1: Prepare + +1. **Write characterization tests** if they don't exist. These capture current behavior — they don't need to be elegant, just comprehensive enough to catch regressions. +2. **Commit** current state. Start from a clean working tree. +3. **Create a branch** for the refactoring. Keep it separate from feature work. + +### Phase 2: Identify + +1. **Smell the code.** Use the smell catalog above to classify what's wrong. +2. **Understand the code.** Read it thoroughly. You must understand what it does before changing it. +3. **Choose the right refactoring.** Pick from the technique catalog. Know what the result looks like before you start. + +### Phase 3: Refactor (Small Steps) + +For each step: +1. **Make one small change.** One refactoring technique at a time. +2. **Compile.** The code should compile after every change. +3. **Run tests.** All tests must pass. If they don't, you've changed behavior. +4. **Commit.** Create a commit with a message like `refactor: extract validateEmail method`. + +Repeat until the smell is resolved. + +### Phase 4: Verify + +1. **All tests pass.** Non-negotiable. +2. **Manual check.** Briefly run the application or review the diff for unintended changes. +3. **Performance.** Ensure no performance regression. Simple refactorings rarely cause them, but check. + +### Phase 5: Clean Up + +1. **Remove stale comments.** If a refactoring made a comment obvious, delete the comment. +2. **Check for dead code.** After refactorings, unused code may emerge. +3. **Final commit.** Summarize the refactoring sequence. + +--- + +## Refactoring Checklist + +### Code Quality +- [ ] Functions are small (< 20 lines preferred, < 50 lines max) +- [ ] Each function does one thing (single responsibility) +- [ ] No duplicated code (DRY) +- [ ] Names describe what, not how +- [ ] No magic numbers or strings +- [ ] Dead code removed + +### Structure +- [ ] Related code is grouped together +- [ ] Module boundaries are clear +- [ ] Dependencies flow in one direction (no cycles) +- [ ] No circular dependencies + +### Conditionals +- [ ] Guard clauses replace deep nesting +- [ ] Complex conditions extracted to named methods +- [ ] Polymorphism replaces type-switching conditionals +- [ ] Null Object pattern where appropriate + +### Type Safety (typed languages) +- [ ] Types defined for all public APIs +- [ ] No `any` usage without qualification +- [ ] Nullable types explicitly marked +- [ ] Type codes replaced with classes/enums + +### Testing +- [ ] Refactored code is tested +- [ ] Edge cases are covered +- [ ] All tests pass after each step +- [ ] Characterization tests capture pre-refactoring behavior + +--- + +## Language-Specific Guidance + +### Java +- Prefer `final` for locals that shouldn't change +- Use IDE automated refactorings (Eclipse/IntelliJ) for mechanical steps +- Leverage the type system: enums, records (Java 14+), sealed classes (Java 17+) + +### JavaScript/TypeScript +- Use destructuring to reduce parameter count +- Prefer `const` over `let` for immutable bindings +- Use TypeScript union types instead of type codes +- Nullish coalescing (`??`) and optional chaining (`?.`) eliminate null-check noise + +### Python +- Use type hints for documenting intent during refactoring +- Use `dataclasses` to replace tuple/data-class patterns +- Use `@property` to replace getters +- Context managers for resource cleanup patterns + +### Go +- Small interfaces preferred: accept interfaces, return structs +- Use named return values when they improve clarity +- Table-driven tests pair well with refactoring +- Avoid deep nesting with early returns + +### Rust +- Use `Result` and `Option` instead of error codes and null +- Pattern matching replaces if-else chains +- `From` trait implementations clean up type conversions +- Derive macros reduce boilerplate + +--- + +## Common Refactoring Sequences + +### Extract Method Sequence +1. Create a new method named after intent +2. Copy code fragment into new method +3. Identify local variables → parameters / return values +4. Call new method from original location +5. Test + +### Replace Conditional with Polymorphism Sequence +1. Create subclasses for each variant +2. Create a factory method that returns the right subclass +3. Move the conditional body to the appropriate subclass method +4. Delete the conditional + +### Extract Class Sequence +1. Identify a coherent subset of fields and methods +2. Create a new class +3. Create an instance from the old class +4. Move fields and methods one at a time +5. Update references in old class +6. Test after each move + +### Inline Class Sequence +1. Identify all callers of the target class +2. Move all methods/fields to the absorbing class +3. Redirect all references to the absorbing class +4. Delete the empty class +5. Test + +--- + +## Safety Protocol + +### Before You Touch Anything +``` +1. Characterization tests → capture what the code does now +2. Git commit → save a known-good state +3. Branch → isolate refactoring from other work +``` + +### Every Single Step +``` +1. One change → one refactoring technique +2. Compile → must compile clean +3. Tests → every test must pass +4. Commit → message: "refactor: " +``` + +### If Tests Break +``` +1. Undo the last change +2. Understand what broke and why +3. Try a smaller step +4. If the test was wrong and behavior was correct, fix the test FIRST, then retry +``` + +### On Completion +``` +1. Full test suite → all tests pass +2. Manual smoke test → quick sanity check +3. Self-review diff → catch unintended changes +4. Final commit → describe the overall transformation +``` + +--- + +## Design Patterns in Refactoring + +### Strategy Pattern +Replace a conditional that chooses an algorithm. **Smell:** Switch on type code with different behavior per branch. **Technique:** Replace Conditional with Polymorphism + Extract Method. + +### Template Method +Extract common algorithm skeleton to superclass, letting subclasses fill in the variants. **Smell:** Duplicate code with slight variations. **Technique:** Form Template Method. + +### State Pattern +Replace a state-based conditional by extracting each state's behavior into a class. **Smell:** Switch on status field with behavior variation. **Technique:** Replace Type Code with State/Strategy. + +### Composite Pattern +Treat individual objects and groups uniformly. **Smell:** Client code has special handling for single vs. collection cases. **Technique:** Extract Interface + Create Composite. + +### Decorator Pattern +Add behavior dynamically by wrapping objects. **Smell:** Conditional logic for optional behaviors. **Technique:** Extract Class + use composition. + +### Null Object Pattern +Replace null checks with a default object. **Smell:** Repeated `if (x == null)` checks. **Technique:** Introduce Null Object. + +--- + +## Edge Cases & Gotchas + +| Scenario | Handling | +|----------|----------| +| No tests exist | Write characterization tests first. Run the code with various inputs, capture outputs. These are your safety net. | +| Refactoring breaks a distant test | FIRST understand why. Maybe the test relied on implementation detail. If so, fix the test to test behavior, not implementation. Then resume. | +| User wants behavior change + refactor together | REFUSE. Do them separately. Refactor first to make the behavior change easy, commit, then change behavior. | +| Method is too complex to step through | Use Replace Method with Method Object. Turn the whole method into a class where each step can be extracted. | +| Refactoring across a large codebase | Extract a micro-service or module boundary first. Then refactor within the boundary. "There is a refactoring for everything except too many refactorings." | +| IDE automated refactoring available | Use it. Modern IDEs can safely rename, extract method, introduce variable, etc. Only do it manually when the IDE can't. | +| Undo needed | `git stash` or `git reset --hard` back to last commit. Small commits make this painless. | + +--- + +## Resources + +- Martin Fowler, *Refactoring: Improving the Design of Existing Code* (2nd Edition, 2018) +- [refactoring.com](https://refactoring.com) — Fowler's online catalog +- [refactoring.guru](https://refactoring.guru) — Illustrated refactoring patterns +- [refactoring.guru/refactoring/catalog](https://refactoring.guru/refactoring/catalog) — full technique catalog (6 categories, 66 techniques) and code-smell taxonomy this skill mirrors \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/review-it/SKILL.md b/pigo/internal/builtinskills/skills/review-it/SKILL.md new file mode 100644 index 0000000..e1506d1 --- /dev/null +++ b/pigo/internal/builtinskills/skills/review-it/SKILL.md @@ -0,0 +1,163 @@ +--- +name: review-it +description: "Code review closeout for Claude Code, Codex, OpenCode, DeepSeek TUI, and Antigravity CLI: local dirty changes, branch vs main, parallel tests." +--- + +# CC Review + +Run automated code review as a closeout check before committing or shipping. Works across multiple AI coding agents. + +Use when: +- user asks for code review / review-it / autoreview +- after non-trivial code edits, before final/commit/ship +- reviewing a local branch or PR branch after fixes + +## Supported Agents + +| Agent | Review Command | Notes | +|-------|---------------|-------| +| Claude Code | `/review` | Built-in, works on uncommitted changes or diff | +| Codex | `codex review` | Pass diff file or let it auto-detect | +| OpenCode | `/review` | Same as Claude Code | +| DeepSeek TUI | `/review` or manual diff review | Pass diff content for analysis | +| Antigravity CLI | `/code-review` | Built-in slash command, auto-detects diff | + +## Contract + +- Treat review output as advisory. Never blindly apply it. +- Verify every finding by reading the real code path and adjacent files. +- Read dependency docs/source/types when the finding depends on external behavior. +- Reject unrealistic edge cases, speculative risks, broad rewrites, and fixes that over-complicate the codebase. +- Prefer small fixes at the right ownership boundary; no refactor unless it clearly improves the bug class. +- Keep going until review returns no accepted/actionable findings. +- If a review-triggered fix changes code, rerun focused tests and rerun review. +- Stop as soon as the review comes back clean with no actionable findings. +- If rejecting a finding as intentional/not worth fixing, add a brief inline code comment only when it explains a real invariant or ownership decision that future reviewers should know. +- Do not push just to review. Push only when the user requested push/ship/PR update. + +## Review Focus + +请 review 当前 diff。不要只看语法和明显 bug,请重点检查以下维度,最后按严重程度排序: + +1. **隐藏副作用 (Hidden Side Effects)** — 变更是否在非显而易见的地方产生级联影响?是否修改了共享状态、全局变量、或外部依赖的行为? +2. **破坏兼容性 (Breaking Compatibility)** — 是否改变了 API 签名、数据结构、配置文件格式、或命令行接口?现有调用方是否会受影响? +3. **边界情况 (Edge Cases)** — null/空值/空集合、极大/极小值、并发/竞态条件、异常路径是否被正确处理? +4. **性能风险 (Performance Risks)** — 是否引入了不必要的循环嵌套、N+1 查询、大对象分配、阻塞 I/O、或锁竞争? +5. **安全风险 (Security Risks)** — 是否存在注入、越权、敏感信息泄露、不安全的反序列化、或依赖版本漏洞? +6. **命名误导 (Naming Misleading)** — 变量/函数/类型名称是否与实际行为不一致?是否存在名不副实或语义模糊的命名? +7. **测试不足 (Insufficient Testing)** — 关键路径、边界条件、错误处理是否缺少测试覆盖?现有测试是否真正验证了期望行为? +8. **未来维护成本 (Future Maintenance Cost)** — 是否引入了不必要的抽象、重复代码、隐式耦合、或难以追踪的控制流?后来者是否容易理解和修改? + +## Pick Target + +### Claude Code / OpenCode / DeepSeek TUI + +Dirty local work (default — `/review` works on uncommitted changes): + +``` +/review +``` + +Branch/PR work — review all changes against base: + +First generate a diff, then review it: + +```bash +git diff origin/main...HEAD > /tmp/review-it.diff +``` + +Then review the diff file with a focused prompt: + +``` +/review the changes in /tmp/review-it.diff against origin/main +``` + +If an open PR exists, use its actual base: + +```bash +base=$(gh pr view --json baseRefName --jq .baseRefName) +git diff "origin/$base"...HEAD > /tmp/review-it.diff +``` + +### Antigravity CLI (`agy`) + +Dirty local work: + +``` +/code-review +``` + +Branch/PR work — review all changes against base: + +```bash +git diff origin/main...HEAD > /tmp/review-it.diff +``` + +Then pass the diff to the review command: + +``` +/code-review the changes in /tmp/review-it.diff against origin/main +``` + +If an open PR exists, use its actual base: + +```bash +base=$(gh pr view --json baseRefName --jq .baseRefName) +git diff "origin/$base"...HEAD > /tmp/review-it.diff +``` + +### Codex + +```bash +# Review uncommitted changes +codex review + +# Review branch diff +git diff origin/main...HEAD > /tmp/review-it.diff +codex review /tmp/review-it.diff +``` + +## Parallel Closeout + +Format first if formatting can change line locations. Then it's OK to run tests and review in parallel: + +```bash +scripts/review-it --parallel-tests "" +``` + +Tradeoff: tests may force code changes that stale the review. If tests or review lead to code edits, rerun the affected tests and rerun review until no accepted/actionable findings remain. + +## Uncommitted vs Branch Review + +Choose the right mode: + +- **Uncommitted changes** (staged/unstaged): use `/review` directly (Antigravity: `/code-review`, Codex: `codex review`) +- **Committed, not pushed**: use `git diff origin/main...HEAD` + review +- **Pushed/PR**: same as committed, against the PR base +- **Clean working tree**: skip review if there's truly nothing to review + +## Helper + +Bundled helper script for parallel test + review orchestration: + +```bash +~/.claude/skills/review-it/scripts/review-it --help +``` + +The helper: +- Detects which agent is running (Claude Code, Antigravity CLI, Codex) via `--agent auto` +- Detects whether to use uncommitted review or branch diff review +- For branch mode: generates diff against `origin/main` (or PR base), then triggers review +- Supports `--parallel-tests` for concurrent test + review execution +- Supports `--dry-run` for checking what command would be used +- Prints `review-it clean: no accepted/actionable findings reported` when review is clean + +## Final Report + +Include: +- review target (uncommitted / branch / PR base) +- tests/proof run +- findings accepted/rejected, briefly why +- the clean review result, or why a remaining finding was consciously rejected + +Do not run another review solely to improve the final report wording. If review exited clean with no actionable findings, report that as clean. \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/review-it/scripts/review-it b/pigo/internal/builtinskills/skills/review-it/scripts/review-it new file mode 100644 index 0000000..ed0d108 --- /dev/null +++ b/pigo/internal/builtinskills/skills/review-it/scripts/review-it @@ -0,0 +1,248 @@ +#!/usr/bin/env bash +# review-it — Code review closeout helper +# Orchestrates /review with optional parallel test execution. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SKILL_DIR="$(dirname "$SCRIPT_DIR")" + +MODE="auto" +BASE="origin/main" +TESTS="" +DRY_RUN=false +AGENT="auto" + +usage() { + cat </dev/null | wc -l | tr -d ' ') +UNSTAGED=$(git diff --name-only 2>/dev/null | wc -l | tr -d ' ') +UNTRACKED=$(git ls-files --others --exclude-standard 2>/dev/null | wc -l | tr -d ' ') +CURRENT_BRANCH=$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "HEAD") +IS_MAIN=false +[[ "$CURRENT_BRANCH" == "main" || "$CURRENT_BRANCH" == "master" ]] && IS_MAIN=true + +HAS_DIRTY=false +[[ "$STAGED" -gt 0 || "$UNSTAGED" -gt 0 || "$UNTRACKED" -gt 0 ]] && HAS_DIRTY=true + +# --- Detect agent --- +detect_agent() { + # Check environment variables / parent process + if [[ -n "${ANTIGRAVITY_CLI:-}" ]] || [[ -n "${GEMINI_CLI:-}" ]]; then + echo "antigravity" + elif [[ -n "${CLAUDE_CODE:-}" ]] || [[ -n "${CLAUDE_CLI:-}" ]]; then + echo "claude" + elif command -v codex &>/dev/null && [[ -n "${CODEX_CLI:-}" ]]; then + echo "codex" + else + # Default to claude as most common + echo "claude" + fi +} + +if [[ "$AGENT" == "auto" ]]; then + DETECTED_AGENT=$(detect_agent) +else + DETECTED_AGENT="$AGENT" +fi + +# --- Resolve review command by agent --- +review_cmd_for() { + local mode="$1" + local diff_file="$2" + case "$DETECTED_AGENT" in + antigravity) + if [[ "$mode" == "local" ]]; then + echo "/code-review" + else + echo "/code-review the changes in $diff_file against $BASE" + fi + ;; + codex) + if [[ "$mode" == "local" ]]; then + echo "codex review" + else + echo "codex review $diff_file" + fi + ;; + claude|*) + if [[ "$mode" == "local" ]]; then + echo "/review" + else + echo "/review the changes in $diff_file against $BASE" + fi + ;; + esac +} + +# --- Resolve mode --- +if [[ "$MODE" == "auto" ]]; then + if $HAS_DIRTY; then + REVIEW_MODE="local" + elif $IS_MAIN; then + echo "review-it: on main with clean tree, nothing to review." + exit 0 + else + REVIEW_MODE="branch" + fi +else + REVIEW_MODE="$MODE" +fi + +# --- Check for PR base --- +if [[ "$REVIEW_MODE" == "branch" ]]; then + if command -v gh &>/dev/null; then + PR_BASE=$(gh pr view --json baseRefName --jq .baseRefName 2>/dev/null || echo "") + if [[ -n "$PR_BASE" ]]; then + BASE="origin/$PR_BASE" + fi + fi + # Fetch the base to ensure we have it + if ! $DRY_RUN; then + git fetch origin "$(echo "$BASE" | sed 's|origin/||')" 2>/dev/null || true + fi +fi + +# --- Build review command --- +DIFF_FILE="" +case "$REVIEW_MODE" in + local) + REVIEW_DESC="uncommitted changes" + ;; + branch) + DIFF_FILE="/tmp/review-it-$$.diff" + REVIEW_DESC="branch $CURRENT_BRANCH vs $BASE" + ;; + *) + echo "review-it: unknown mode $REVIEW_MODE" + exit 1 + ;; +esac + +REVIEW_CMD=$(review_cmd_for "$REVIEW_MODE" "$DIFF_FILE") + +# --- Dry run --- +if $DRY_RUN; then + echo "mode: $REVIEW_MODE" + echo "agent: $DETECTED_AGENT" + echo "target: $REVIEW_DESC" + echo "review: $REVIEW_CMD" + echo "dirty: staged=$STAGED unstaged=$UNSTAGED untracked=$UNTRACKED" + if [[ "$REVIEW_MODE" == "branch" ]]; then + echo "diff: $DIFF_FILE" + echo "base: $BASE" + fi + if [[ -n "$TESTS" ]]; then + echo "tests: $TESTS" + fi + exit 0 +fi + +# --- Run --- +echo "review-it: mode=$REVIEW_MODE agent=$DETECTED_AGENT target=$REVIEW_DESC" +echo "" + +# Generate diff for branch mode +if [[ "$REVIEW_MODE" == "branch" ]]; then + if ! git diff "$BASE"...HEAD > "$DIFF_FILE" 2>/dev/null; then + echo "review-it: failed to generate diff against $BASE" + exit 1 + fi + if [[ ! -s "$DIFF_FILE" ]]; then + echo "review-it: no changes between $BASE and HEAD" + rm -f "$DIFF_FILE" + exit 0 + fi + echo "review-it: generated diff ($(wc -l < "$DIFF_FILE" | tr -d ' ') lines) at $DIFF_FILE" + echo "" +fi + +# Run tests in background if requested +TEST_PID="" +if [[ -n "$TESTS" ]]; then + echo "review-it: running tests in background: $TESTS" + eval "$TESTS" > /tmp/review-it-tests-$$.log 2>&1 & + TEST_PID=$! + echo "review-it: test PID=$TEST_PID" + echo "" +fi + +# Run review (this is a signal to the user — the actual /review +# is executed by Claude Code when it reads this output) +cat </dev/null || true + TEST_EXIT=$? + echo "" + echo "--- test output ---" + cat /tmp/review-it-tests-$$.log + echo "--- end test output ---" + rm -f /tmp/review-it-tests-$$.log + if [[ $TEST_EXIT -ne 0 ]]; then + echo "" + echo "review-it: tests FAILED (exit=$TEST_EXIT)" + exit $TEST_EXIT + else + echo "" + echo "review-it: tests passed" + fi +fi + +# Cleanup +if [[ "$REVIEW_MODE" == "branch" ]]; then + rm -f "$DIFF_FILE" +fi + +echo "" +echo "review-it: done" \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/ship-it/SKILL.md b/pigo/internal/builtinskills/skills/ship-it/SKILL.md new file mode 100644 index 0000000..ecbc7b9 --- /dev/null +++ b/pigo/internal/builtinskills/skills/ship-it/SKILL.md @@ -0,0 +1,181 @@ +--- +name: ship-it +description: Code commit, PR creation, merge, and issue closure workflow via GitHub CLI (gh). Triggers after a goal (GitHub Issue) implementation is complete — commit code, push branch, create PR, merge, then close the issue. Use when the user says "提交代码", "commit and merge", "创建PR", "合入", "关闭issue", "ship-it", or when a goal implementation is done and code needs to be shipped. +allowed-tools: + - Bash(git:*) + - Bash(gh:*) +--- + +# After-Goal: 代码提交、PR 合入、Issue 关闭工作流(GitHub) + +完成 GitHub Issue 实现后的标准收尾流程:提交代码 → 推送分支 → 创建 PR → 合入 → 关闭 Issue。 + +## 前置条件 + +- 当前 git 仓库有已实现的代码变更 +- 已知 Issue 编号(如 `#42`) +- gh CLI 已登录(`gh auth status` 可验证) + +## 工作流 + +### Step 1: 提交代码 + +```bash +# 1a. 检查变更状态 +git status +git diff --stat HEAD + +# 1b. 暂存本次 Issue 相关的文件(不要 add 不相关的文件) +git add + +# 1c. 提交,commit message 关联 Issue +git commit -m "$(cat <<'EOF' +{简要描述} (#issue-number) + +{可选的详细说明} +EOF +)" +``` + +**关键规则:** +- commit message 中包含 `#issue-number` 以关联 Issue +- 只暂存当前 Issue 相关的文件,不要混入其他变更 + +### Step 2: 推送分支 + +```bash +# 如果还在 main/master 上,先创建功能分支 +git checkout -b {branch-name} # 如已在功能分支则跳过 + +# 推送到远程 +git push -u origin {branch-name} +``` + +分支命名建议:`feat/issue-42-short-desc` 或 `fix/issue-42-short-desc` + +### Step 3: 创建 PR + +```bash +gh pr create \ + --title "{简要描述}" \ + --body "$(cat <<'EOF' +## Summary +- 实现内容概述 + +Closes #{issue-number} + +## Test plan +- [ ] 测试项 1 +- [ ] 测试项 2 +EOF +)" +``` + +**关键规则:** +- PR body 中写 `Closes #N` 或 `Fixes #N`,合入后 GitHub 自动关闭 Issue +- title 简洁,不超过 70 字符 + +### Step 4: 合入 PR + +```bash +# 4a. 查看 PR 状态(确认 checks 通过) +gh pr checks + +# 4b. 合入(默认 merge commit,可选 --squash 或 --rebase) +gh pr merge --squash --delete-branch +``` + +**参数说明:** +- `--squash`: 压缩为单个 commit 合入(推荐) +- `--rebase`: rebase 合入 +- `--merge`: 普通 merge commit +- `--delete-branch`: 合入后删除远程分支 + +### Step 5: 添加实现总结评论 + +PR 合入后,始终在 Issue 上添加实现总结评论,方便后续直接从 Issue 回溯代码变更。 + +```bash +gh issue comment {issue-number} --body "$(cat <<'EOF' +## 实现总结 +- **核心变更**:{从 PR body 提取的实现摘要} +- **PR**: #{pr-number} +- **Commit**: {hash} +EOF +)" +``` + +**关键规则:** +- 无论是 auto-close 还是手动 close,都必须添加此评论 +- 评论内容从 PR body 的 Summary 部分提取,保持简洁(3-5 条 bullet) +- 附加 PR 编号和 commit hash,方便直接跳转 + +### Step 6: 手动关闭 Issue(仅当未自动关闭时) + +如果 PR body 中已写 `Closes #N`,合入后 Issue 会自动关闭,跳过此步。否则手动关闭: + +```bash +gh issue close {issue-number} --reason completed +``` + +## 错误处理 + +| 场景 | 处理方式 | +|------|---------| +| `gh pr checks` 有失败项 | 查看失败原因,修复后追加 commit 推送 | +| PR 有 merge conflict | `git fetch origin main && git rebase origin/main`,解决冲突后 force push | +| `gh pr merge` 被 branch protection 阻止 | 确认 required reviews 已满足,或请 reviewer approve | +| Issue 合入后未自动关闭 | 确认 PR body 包含 `Closes #N`,或执行 Step 6 手动 `gh issue close` | + +## 完整示例 + +```bash +# 创建分支并提交 +git checkout -b feat/issue-42-case-model +git add cases/case.go cases/case_test.go +git commit -m "$(cat <<'EOF' +Add Case data model and Markdown read/write (#42) + +Define Case struct with YAML frontmatter + Markdown body +serialization. Provide WriteCase/ReadCase/ListCases/UpdateCase. +EOF +)" + +# 推送 +git push -u origin feat/issue-42-case-model + +# 创建 PR +gh pr create \ + --title "Add Case data model and Markdown read/write" \ + --body "$(cat <<'EOF' +## Summary +- Define Case struct with YAML frontmatter + Markdown body +- Implement WriteCase/ReadCase/ListCases/UpdateCase functions +- Add comprehensive test coverage + +Closes #42 + +## Test plan +- [x] Unit tests pass +- [x] go vet / lint clean +EOF +)" + +# 确认 checks 通过后合入 +gh pr checks +gh pr merge --squash --delete-branch + +# 添加实现总结评论 +gh issue comment 42 --body "$(cat <<'EOF' +## 实现总结 +- **核心变更**:Define Case struct with YAML frontmatter + Markdown body +- **核心变更**:Implement WriteCase/ReadCase/ListCases/UpdateCase +- **PR**: #43 +- **Commit**: abc1234 +EOF +)" + +# 切回主分支 +git checkout main +git pull +``` diff --git a/pigo/internal/builtinskills/skills/smell/README.md b/pigo/internal/builtinskills/skills/smell/README.md new file mode 100644 index 0000000..d888e81 --- /dev/null +++ b/pigo/internal/builtinskills/skills/smell/README.md @@ -0,0 +1,39 @@ +# Smell — Architecture Bad Smell Detector + +Analyze a codebase to find violations of software architecture principles, anti-patterns, and code "bad smells." Produces a comprehensive, actionable markdown report. + +## Features + +- Scans project structure for architectural anti-patterns (Big Ball of Mud, Distributed Monolith, etc.) +- Detects coupling and cohesion issues (God Objects, circular dependencies, feature envy) +- Identifies design principle violations (SOLID, DRY, KISS, YAGNI) +- Finds code-level smells (Long Method, Primitive Obsession, Magic Numbers) +- Assesses testing health (missing tests, test-implementation coupling) +- Outputs a structured markdown report with severity levels and refactoring roadmap + +## Knowledge Base + +Built on architectural knowledge from: +- [awesome-software-architecture](https://github.com/mehdihadeli/awesome-software-architecture) +- Big Ball of Mud (Foote & Yoder, 1997) +- Clean Architecture, Onion Architecture, Hexagonal Architecture +- Domain-Driven Design, CQRS, Event-Driven Architecture +- SOLID, DRY, KISS, YAGNI, GRASP principles + +## Usage + +Trigger with prompts like: + +- "smell" or "/smell" +- "find code smells" +- "detect architecture anti-patterns" +- "analyze architecture quality" +- "找出坏味道" +- "架构坏味道" +- "代码坏味道检测" +- "反模式分析" + +## Files + +- `SKILL.md` — Skill definition and comprehensive anti-pattern knowledge base +- `test-prompts.json` — Test prompts for validation \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/smell/SKILL.md b/pigo/internal/builtinskills/skills/smell/SKILL.md new file mode 100644 index 0000000..70dce39 --- /dev/null +++ b/pigo/internal/builtinskills/skills/smell/SKILL.md @@ -0,0 +1,690 @@ +--- +name: smell +description: "Detect software architecture bad smells, algorithmic complexity hotspots, and anti-patterns in a codebase. Produces a detailed markdown report identifying violations of architectural principles, design patterns, code quality, and performance complexity. Triggers on: smell, code smell, architecture smell, find anti-patterns, detect bad smells, complexity analysis, 代码坏味道, 架构坏味道, 反模式, 找出坏味道, 复杂度分析." +user-invocable: true +--- + +# Smell — Architecture Bad Smell Detector + +Analyze a codebase to find violations of software architecture principles, anti-patterns, code "bad smells," and algorithmic complexity hotspots. Produce a comprehensive, actionable markdown report. + +**Knowledge base:** This skill encodes architectural patterns, anti-patterns, code smells, and algorithmic complexity heuristics drawn from industry research and practice, including the classic code smells catalog by Martin Fowler / Kent Beck (as organized on refactoring.guru: Bloaters, Object-Orientation Abusers, Change Preventers, Dispensables, Couplers). + +--- + +## The Job + +1. Understand the scope — ask what part of the project to analyze (full project, specific module, or recent changes) +2. Scan the codebase using `find`, `grep`, and `Agent` (Explore subagent) to gather evidence +3. Identify architectural smells and anti-patterns +4. Generate a detailed markdown report saved to `tasks/smell-report-[timestamp].md` +5. Present a summary of findings to the user + +--- + +## Step 1: Scope Clarification + +Ask the user: + +``` +What scope should I analyze? + A. Entire project (thorough, may take time) + B. Specific module/directory: [please specify] + C. Only recently changed files (git diff) + D. Only architectural-level issues (skip low-level code smells) +``` + +If the user doesn't specify, default to option A for small projects (< 100 files) or C for large projects. + +--- + +## Step 2: Evidence Gathering + +**Use the Explore subagent** (`Agent` with `subagent_type: "Explore"`) to scan the codebase for architectural patterns and anti-patterns. Run multiple parallel explorations: + +### Exploration Commands + +Run these in parallel to gather evidence efficiently: + +1. **Project Structure Scan:** Map the directory tree, identify the architectural style (layered, modular monolith, microservices, etc.) +2. **Dependency Analysis:** Find import/include patterns, check for circular dependencies, identify coupling hotspots +3. **Module/Component Scan:** Identify God Objects (files > 500 lines), check cohesion, check single responsibility violations +4. **Pattern Detection:** Look for known anti-pattern signatures (static cling, service locator abuse, leaky abstractions) +5. **Testing Scan:** Check test coverage patterns, test file locations, test-to-code ratios +6. **Naming & Clarity Scan:** Flag misleading names, overly generic names (Manager, Helper, Util), inconsistent naming conventions +7. **Complexity Scan:** Detect algorithmic complexity hotspots — nested loops, N+1 queries, repeated scans, sort-in-loop, expensive recomputation in render paths + +### Key Heuristics + +| Category | Smell | Detection Heuristic | +|----------|-------|-------------------| +| **Architecture** | Big Ball of Mud | No clear directory structure; everything in root or one flat folder; no separation of concerns | +| **Architecture** | Violated Layer Boundaries | Inner layers importing outer layers; infrastructure code in domain/core layer | +| **Architecture** | Missing Architecture | No `src/`, `lib/`, `core/` separation; SQL inline with UI code; HTTP handlers mixed with business logic | +| **Architecture** | Distributed Monolith | Microservices sharing a database; services that can't deploy independently | +| **Architecture** | Anemic Domain Model | Model/entity classes with only getters/setters and no behavior; all logic in services | +| **Architecture** | CQRS Without Need | Separate read/write models for simple CRUD; unnecessary complexity | +| **Architecture** | Over-Layered Architecture | Excessive layers/tiers that add pass-through code with no real value | +| **Architecture** | Over-Abstraction | So many indirections/interfaces/generics that you get lost following the code | +| **Architecture** | Futuristic Architecture | Speculative flexibility for requirements that may never come (predicting the future) | +| **Architecture** | Technology-Enthusiast Architecture | Shiny/unproven tech adopted in production because it's new, not because it fits | +| **Architecture** | Overkill Architecture | Heavyweight architecture/tech thrown at a simple problem | +| **Architecture** | Cloud/Visio Architecture | Diagrams disconnected from the actual code and runtime reality | +| **Coupling** | Circular Dependencies | Module A imports B, B imports A; detected via import graph analysis | +| **Coupling** | Content Coupling | One module directly accesses another's internal/private members | +| **Coupling** | Common Coupling | Excessive global variables/shared mutable state; singleton abuse | +| **Coupling** | Stamp Coupling | Passing large data structures when only a few fields are needed | +| **Cohesion** | God Object | Single class/module > 500 lines; > 20 public methods; handles unrelated concerns | +| **Cohesion** | Shotgun Surgery | A single change requires touching 5+ files across unrelated modules | +| **Cohesion** | Feature Envy | Method calls foreign class methods more than its own class methods | +| **Cohesion** | Data Clumps | Same group of 3+ parameters appearing together in multiple method signatures | +| **Design** | Leaky Abstractions | Implementation details (DB queries, HTTP calls) exposed through interfaces | +| **Design** | Static Cling | Excessive use of static methods; static state that prevents testability | +| **Design** | Service Locator Abuse | DI container passed around instead of proper constructor injection | +| **Design** | Violated SOLID | SRP violations, OCP violations (switch/if-else chains on types), ISP violations (fat interfaces) | +| **Design** | Switch Statements | Same `switch`/if-else chain on a type code appearing in multiple places; should be polymorphism | +| **Design** | Refused Bequest | Subclass inherits methods/fields it doesn't use or overrides them to throw/no-op | +| **Design** | Alternative Classes w/ Different Interfaces | Two classes do the same thing but have differently-named methods | +| **Design** | Parallel Inheritance Hierarchies | Creating a subclass in one hierarchy forces a matching subclass in another | +| **Design** | Speculative Generality | Unused abstract classes, hooks, params, or generics "for future needs" (YAGNI) | +| **Design** | Incomplete Library Class | Wrapping/patching a third-party class because it lacks needed methods | +| **Cohesion** | Divergent Change | One module changed for many unrelated reasons (opposite of Shotgun Surgery) | +| **Cohesion** | Data Class | Class with only fields + getters/setters, no behavior (anemic data bag) | +| **Cohesion** | Lazy Class | Class/module that does too little to justify its existence | +| **Coupling** | Inappropriate Intimacy | Two classes access each other's private/internal parts too much | +| **Coupling** | Message Chains | Long call chains `a.getB().getC().getD()` (Law of Demeter violation) | +| **Coupling** | Middle Man | Class that only delegates every call to another class | +| **Code** | Temporary Field | Instance field set/used only in certain circumstances, empty otherwise | +| **Code** | Duplicated Code | Identical/similar logic appearing in 3+ places; copy-paste patterns | +| **Code** | Long Method | Methods > 50 lines; deep nesting (> 3 levels) | +| **Code** | Long Parameter List | Methods with > 4 parameters | +| **Code** | Primitive Obsession | Using strings/ints instead of domain types (e.g., `string email` instead of `Email` type) | +| **Code** | Magic Numbers/Strings | Hardcoded literals without named constants | +| **Code** | Comments as Deodorant | Excessive comments explaining bad code instead of refactoring | +| **Code** | Dead Code | Unused imports, unreachable code, commented-out blocks | +| **Testing** | No Tests | Modules with zero test coverage | +| **Testing** | Test-Implementation Coupling | Tests that assert internal implementation details instead of behavior | +| **Testing** | Slow Tests | Tests doing real I/O, database calls, network requests without mocking | +| **Naming** | Vague Names | `Manager`, `Handler`, `Processor`, `Helper`, `Util`, `Service`, `Data`, `Info` used excessively without context | +| **Naming** | Inconsistent Naming | Snake_case and camelCase mixed; different patterns for same concept | +| **Readability** | Deep Nesting (Arrow Anti-Pattern) | Loops/conditionals nested > 3 levels deep; rightward-drifting "arrow" shape hard to trace | +| **Complexity** | Nested Loops (O(n^2)+) | Loop inside loop; forEach inside for; map inside map; nested iteration suggesting polynomial complexity | +| **Complexity** | Repeated Linear Scan | `includes()`/`indexOf()`/`.find()` inside a loop; O(n*m) membership check on list instead of Set/Map | +| **Complexity** | Sort-in-Loop | `.sort()` or `sorted()` called inside iterative code; repeated O(n log n) when sort-once suffices | +| **Complexity** | N+1 Query Pattern | Database/API/HTTP call inside a loop; `fetch`/`query`/`execute`/`findMany` per iteration instead of batch | +| **Complexity** | Render-Path Recompute | `.filter().map().sort()` chains in component render body; expensive transforms without memoization | +| **Complexity** | Pairwise Comparison | Nested iteration comparing every element with every other; O(n^2) when sort+two-pointer would be O(n log n) | +| **Complexity** | Unnecessary Recompute | Same expensive computation repeated without caching; missing `useMemo`/`memo`/lazy eval | +| **Complexity** | Wrong Data Structure | Array used where Set/Map would give O(1) lookup; List where Queue/Heap/Stack is natural fit | + +--- + +## Step 3: Report Generation + +Generate the report in this structure: + +```markdown +# Architecture Smell Report + +**Project:** [project-name] +**Scope:** [scope description] +**Date:** [date] +**Analyzer:** smell skill (Ducc) + +--- + +## Executive Summary + +[2-3 paragraph summary: architectural style detected, overall health assessment, and top 3-5 critical issues] + +--- + +## Architectural Style Detected + +[Identify the architectural style: Layered, Modular Monolith, Microservices, Hexagonal, Clean Architecture, or Big Ball of Mud] + +### Style Expectations vs. Reality + +| Expectation | Reality | Status | +|-------------|---------|--------| +| [e.g., Clear layer separation] | [what was found] | ✅/⚠️/🔴 | + +--- + +## Findings by Category + +### 🔴 Critical Issues (Must Fix) + +[Issues that fundamentally undermine architecture] + +### 🟡 Warnings (Should Fix) + +[Issues that degrade maintainability but don't block function] + +### 🔵 Suggestions (Nice to Fix) + +[Minor improvements that would increase quality] + +--- + +## Detailed Findings + +### Finding #1: [Title] + +- **Category:** [Architecture/Coupling/Cohesion/Design/Code/Testing/Naming/Complexity] +- **Severity:** 🔴 Critical / 🟡 Warning / 🔵 Suggestion +- **Anti-Pattern:** [Name of anti-pattern] +- **Location:** [file:line references] +- **Principle Violated:** [SOLID/DRY/KISS/etc.] +- **Description:** [What was found and why it's a problem] +- **Evidence:** [Code snippet or structure description] +- **Recommendation:** [Specific fix, with refactoring approach] + +--- + +## Dependency Graph Analysis + +[Summary of module dependencies, circular dependencies found, coupling hotspots] + +--- + +## Module Health Scorecard + +| Module | Lines | God Object Risk | Coupling | Cohesion | Test Coverage | Health | +|--------|-------|----------------|----------|----------|---------------|--------| +| [name] | [N] | [Low/Med/High] | [Low/Med/High] | [Low/Med/High] | [% or N/A] | 🟢/🟡/🔴 | + +--- + +## Smell Distribution + +| Category | Count | Critical | Warning | Suggestion | +|----------|-------|----------|---------|------------| +| Architecture | [N] | [N] | [N] | [N] | +| Coupling | [N] | [N] | [N] | [N] | +| Cohesion | [N] | [N] | [N] | [N] | +| Design | [N] | [N] | [N] | [N] | +| Code | [N] | [N] | [N] | [N] | +| Testing | [N] | [N] | [N] | [N] | +| Naming | [N] | [N] | [N] | [N] | +| Complexity | [N] | [N] | [N] | [N] | + +--- + +## Refactoring Roadmap + +### Immediate Actions (This Sprint) +1. [Actionable fix 1] +2. [Actionable fix 2] + +### Short-Term (1-3 Months) +1. [Structural improvement 1] +2. [Structural improvement 2] + +### Long-Term (3-12 Months) +1. [Architectural transformation 1] +2. [Architectural transformation 2] + +--- + +## Appendix: Anti-Pattern Reference + +[A condensed reference of anti-patterns checked, with brief descriptions] +``` + +--- + +## Step 4: Save and Present + +Save the report to `tasks/smell-report-[YYYY-MM-DD-HHmm].md` and present a brief summary to the user. + +--- + +## Anti-Pattern Knowledge Base + +This section documents the architectural anti-patterns and bad smells the skill knows about. + +### Architectural Anti-Patterns + +#### Big Ball of Mud +The most common de-facto architecture. A haphazardly structured, sprawling system with no perceivable architecture. Characterized by: +- Promiscuous sharing of information between distant elements +- Global or duplicated important state +- Structure eroded beyond recognition or never defined +- Repeated expedient repair ("duct tape and bailing wire") +- Forces: Time pressure, cost, inexperience, complexity, change, scale +- **Remedy:** Define architecture boundaries, refactor incrementally, apply SHEARING LAYERS, KEEP IT WORKING + +#### Distributed Monolith +Microservices that must be deployed together. Symptoms: +- Services share a database +- Synchronous chains of service calls +- Changes require coordinated deployments +- **Remedy:** Decouple data stores, introduce async messaging, enforce bounded contexts + +#### Anemic Domain Model +Domain objects with only getters/setters (data bags), all logic in services. Violates: +- "Tell, Don't Ask" principle +- Rich Domain Model pattern from DDD +- **Remedy:** Move behavior into domain objects, use domain services only for cross-aggregate operations + +#### God Object +A class that knows too much or does too much. Characteristics: +- > 500 lines or > 20 public methods +- Handles unrelated concerns +- Difficult to test in isolation +- Single Responsibility Principle violation +- **Remedy:** Extract cohesive groups of methods into dedicated classes + +#### Leaky Abstractions +Abstractions that expose implementation details. Signs: +- Interface methods named after implementation (e.g., `SaveToPostgres`, `FetchFromRedis`) +- Consumers catching implementation-specific exceptions +- Configuration details exposed through abstractions +- **Remedy:** Design interfaces from the consumer's perspective, hide implementation details + +#### Static Cling +Excessive use of static methods/state. Problems: +- Untestable (can't mock static calls) +- Hidden dependencies +- Thread-safety issues with static state +- **Remedy:** Use dependency injection, convert stateless statics to instance methods + +#### Service Locator Abuse +Using a service locator instead of dependency injection. Issues: +- Hidden dependencies (dependencies not visible in constructor) +- Runtime errors instead of compile-time errors +- Testing difficulty +- **Remedy:** Use constructor injection, register dependencies at composition root + +#### Violated Layer Boundaries (Clean/Onion/Hexagonal Architecture) +In layered architectures: +- **Clean Architecture:** Outer layers (frameworks) leaking into inner layers (use cases, entities) +- **Onion Architecture:** Infrastructure concerns in domain core +- **Hexagonal Architecture:** Business logic coupled to specific adapters instead of ports +- **Remedy:** Apply dependency inversion, define clear port interfaces + +#### CQRS Overuse +Applying CQRS to simple CRUD. Signs: +- Separate read/write models for trivial data access +- Event sourcing when events don't add business value +- Unnecessary complexity +- **Remedy:** Use CQRS only when read/write models genuinely differ or have different scaling needs + +#### Vertical Slice Contamination +In Vertical Slice Architecture: +- Cross-slice coupling (one feature directly calling another) +- Shared service classes undermining slice independence +- **Remedy:** Use events/messages for cross-slice communication, duplicate simple logic if needed + +### Top Ten Software Architecture Mistakes + +A set of architecture-level anti-patterns describing over- and under-engineering. The common thread: **architecture disconnected from real needs and reality.** The opposite extreme (too little architecture) is equally a smell. + +#### Over-Layered / Multitier Architecture +"Layers on layers on layers." Adding tiers beyond what the problem needs: +- Each layer just forwards calls to the next with no transformation or value +- Simple read requires touching 6+ classes across 4 layers +- **Remedy:** Collapse pass-through layers; keep only layers that carry real responsibility + +#### Over-Abstraction +Abstraction piled on until the code is impossible to follow: +- Excessive interfaces, generics, factories, and indirection for single implementations +- You can't tell what actually runs without stepping through many hops +- **Remedy:** Inline single-implementation abstractions; abstract only at real variation points (rule of three) + +#### Futuristic Architecture +Solution built for imagined future requirements that no one can actually predict: +- Extensibility points, plugin systems, config knobs nothing uses +- Most speculative flexibility is wasted effort — closely related to Speculative Generality and YAGNI +- **Remedy:** Build for today's known requirements; add flexibility when a real second case arrives + +#### Technology-Enthusiast Architecture +New/shiny technology put into production because the architect liked it: +- Unproven tech adopted without validating it fits the problem or scales +- Chasing trends over stability +- **Remedy:** Evaluate tech against actual requirements; prefer proven tools; prototype before committing + +#### Overkill Architecture +A simple problem solved with a disproportionate amount of architecture and technology: +- Microservices, event sourcing, k8s for a CRUD app with a handful of users +- **Remedy:** Match architecture weight to problem size (KISS); start simple, evolve when justified + +#### Cloud / Visio Architecture +"Architecture" that exists only in nice diagrams, disconnected from the code and runtime reality: +- Diagrams don't match what's actually deployed; boxes and arrows with no code correspondence +- **Remedy:** Keep architecture docs grounded in and verified against the real system + +> **Note on the opposite extreme:** total *lack* of architecture (no boundaries, no structure) is equally a smell — see [Big Ball of Mud](#big-ball-of-mud) and Missing Architecture. Both under- and over-engineering are failures. + +### Coupling & Cohesion Smells + +#### Circular Dependencies +Module A → Module B → Module A. Detected via: +- Import graph analysis +- "Cannot access before initialization" errors +- **Remedy:** Extract shared interface/common module, apply dependency inversion + +#### Content Coupling +One module directly modifying another's internal state. Signs: +- Direct field access across module boundaries +- `friend`/package-private abuse +- **Remedy:** Use public APIs, encapsulate internal state + +#### Common Coupling (Global State) +Multiple modules depending on shared global mutable state: +- Global variables, singletons with mutable state +- Ambient context (e.g., `CurrentUser` static property) +- **Remedy:** Parameterize, use dependency injection, make state explicit + +#### Stamp Coupling +Passing entire data structures when only a few fields needed: +- Functions receiving large DTOs but using one field +- **Remedy:** Create focused parameters or smaller interfaces (ISP) + +#### Shotgun Surgery +A single change requires modifications across many files: +- Adding a field touches 5+ files in different modules +- **Remedy:** Consolidate related behavior, apply Single Responsibility + +#### Feature Envy +A method that uses another class's methods more than its own: +- Method calls `other.foo()`, `other.bar()`, `other.baz()` with few self-calls +- **Remedy:** Move the method to the class it envies + +#### Data Clumps +Same group of fields appearing together in multiple places: +- `(street, city, zip)` appearing in 5 method signatures +- **Remedy:** Extract into a value object + +#### Divergent Change +One module/class is repeatedly changed for many *unrelated* reasons (the opposite of Shotgun Surgery): +- "I always change these three methods for DB changes, and those two for UI changes" in the same class +- **Remedy:** Split the class along its axes of change (Single Responsibility) + +#### Inappropriate Intimacy +Two classes are too entangled with each other's internals: +- Reaching into another class's private fields, tight bidirectional references +- **Remedy:** Move methods/fields to the class they belong to, extract a shared class, or replace with delegation + +#### Message Chains +Long navigation chains like `a.getB().getC().getD().doThing()`: +- Client coupled to the whole object graph; violates the Law of Demeter +- **Remedy:** Hide delegation — add a method on the first object that returns what the client needs + +#### Middle Man +A class that delegates almost all of its work to another class: +- Most methods just forward calls; adds indirection without value +- **Remedy:** Remove the middle man and let clients talk to the real object (inline the delegation) + +#### Parallel Inheritance Hierarchies +Every time you add a subclass to one hierarchy, you must add one to another: +- `Shape`/`ShapeRenderer`, `Employee`/`EmployeePermission` growing in lockstep +- **Remedy:** Merge hierarchies or make one hierarchy reference the other instead of mirroring it + +### Code-Level Smells + +#### Long Method +- Methods > 50 lines (or whatever suits the language) +- Deep nesting > 3 levels +- Multiple levels of abstraction mixed +- **Remedy:** Extract methods at same abstraction level, compose + +#### Long Parameter List +- Methods with > 4 parameters +- Boolean flags controlling behavior +- **Remedy:** Introduce parameter object, split method, remove flag arguments + +#### Duplicated Code +- Identical or near-identical logic in 3+ places +- Copy-paste with slight variations +- **Remedy:** Extract shared method, apply Template Method or Strategy pattern + +#### Primitive Obsession +Using primitives instead of domain types: +- `string` for Email, PhoneNumber, URL +- `int` for Money, Age, Quantity +- `decimal` without Currency context +- **Remedy:** Create value objects with validation and behavior + +#### Magic Numbers/Strings +- Hardcoded literals without explanation +- `if (status == 3)` instead of `if (status == Status.COMPLETED)` +- **Remedy:** Extract named constants or enums + +#### Comments as Deodorant +- Comments that explain what code does (code should be self-documenting) +- Commented-out code blocks +- "TODO" comments accumulating without resolution +- **Remedy:** Refactor to make code clear, delete dead code, track TODOs as issues + +#### Deep Nesting (Arrow Anti-Pattern) +Loops and conditionals nested so deeply the code drifts rightward into an "arrow" shape: +- `if { if { for { if { ... } } } }` — hard to trace which conditions hold at any point +- Usually > 3 levels of indentation in one function +- **Remedy:** Guard clauses / early returns, extract nested blocks into methods, invert conditions, replace conditional with polymorphism + +#### Dead Code +- Unused imports, variables, functions +- Unreachable branches +- Commented-out code in version control +- **Remedy:** Delete it (git history preserves it if needed) + +#### Data Class +A class that is only fields plus getters/setters, with no meaningful behavior: +- A "data bag" other classes reach into and manipulate from outside +- Closely related to Anemic Domain Model at the class level +- **Remedy:** Move the behavior that operates on the data into the class ("Tell, Don't Ask") + +#### Lazy Class +A class/module that no longer does enough to justify its existence: +- Left over after refactoring, or an abstraction that never grew +- **Remedy:** Inline it into its caller or collapse the hierarchy + +#### Speculative Generality +Abstractions, hooks, parameters, or generics added for hypothetical future needs: +- Unused abstract base classes, unused parameters, "just in case" configuration +- Violates YAGNI +- **Remedy:** Remove unused abstraction; add it when a real second use case appears + +#### Temporary Field +An instance field that is only set/used in certain circumstances and empty otherwise: +- Fields populated only during one algorithm, confusing readers the rest of the time +- **Remedy:** Extract the field + the methods that use it into their own class (Extract Class / introduce a Method Object) + +### Testing Smells + +#### No Tests +- Modules with zero test coverage +- Business logic without unit tests +- **Remedy:** Write characterization tests first, then add behavior tests + +#### Test-Implementation Coupling +- Tests asserting internal method calls, private state, or implementation details +- Tests breaking on refactoring without behavior changes +- **Remedy:** Test through public APIs, assert behavior not implementation + +#### Test Environment Dependency +- Tests depending on file system, network, database, system clock without mocking +- Non-deterministic tests (flaky tests) +- **Remedy:** Use test doubles, control environment, use DI + +### Complexity Smells (Algorithmic Anti-Patterns) + +Complexity smells indicate code whose runtime grows inefficiently with input size. These are not mere "micro-optimizations" — they are algorithmic choices that cause real performance degradation at scale. + +#### Nested Loops (O(n^2) and Worse) +Two or more loops nested inside each other, producing polynomial complexity. +- **Detection:** `for`/`while` inside another `for`/`while`; `forEach`/`map` inside `forEach`/`map`; loop containing another loop (any depth) +- **Impact:** O(n^2) for double-nested, O(n^3) for triple; explodes with moderate data sizes +- **Remedy:** + - Build a Map/Set index for the inner collection → O(n+m) + - Sort + two-pointer approach → O(n log n) + - Group/bucket data before iterating + - Sweep-line for interval/range problems +- **Correctness checks:** Does order matter? Are there duplicate keys? Is the original picking first/last/all matches? + +#### N+1 Query Pattern +A database query, API call, or I/O operation inside a loop body. +- **Detection:** `fetch()`/`axios()`/`query()`/`execute()`/`findMany()`/`findOne()`/`findUnique()`/`select()`/`where()` inside any loop construct +- **Impact:** 1 + N round-trips instead of 1; network latency multiplied by item count +- **Remedy:** + - Batch fetch by IDs: `SELECT * FROM x WHERE id IN (...)` then join in memory + - Use ORM eager-loading / `include` / `preload` / DataLoader + - Bulk API endpoints accepting arrays + - Preserve: auth filters, tenancy isolation, ordering, pagination, error semantics +- **Correctness checks:** Don't fetch records the original per-item logic wouldn't authorize; preserve missing-record behavior + +#### Repeated Linear Scan (Missing Index) +Linear search (`includes`, `indexOf`, `.find`, `in_array`) inside a loop, where a Set/Map would give O(1) lookup. +- **Detection:** `.includes()` / `.indexOf()` / `.find()` / `.findIndex()` / `in_array()` / `contains()` inside a loop body +- **Impact:** O(n*m) instead of O(n+m) — each iteration scans the entire collection +- **Remedy:** Build a `Set` (for membership) or `Map` (for key→value lookup) once before the loop +- **Correctness checks:** Does equality semantics change after Set conversion? JavaScript object identity vs. value equality; Python hashability + +#### Sort-in-Loop +Sorting inside a loop body, repeating O(n log n) work unnecessarily. +- **Detection:** `.sort()` / `sorted()` / `sort()` inside any iterative block +- **Impact:** O(k * n log n) instead of O(n log n) — sort repeated k times +- **Remedy:** + - Sort once outside the loop + - Maintain a heap (PriorityQueue) if incremental top-K is needed + - Use binary search/insertion into sorted collection +- **Correctness checks:** Is each intermediate sorted state externally observable? Does comparator depend on loop-local state? + +#### Render-Path Recompute (UI Complexity) +Expensive data transformation (filter→map→sort chains) inside UI component render bodies, recomputed on every render. +- **Detection:** `.filter().map().sort().reduce()` chains inside React/Vue/Svelte component function bodies; inside `function Component()` or `const Component = () =>` in JSX/TSX +- **Impact:** Re-derivation on every state change even if inputs unchanged; jank with large collections +- **Remedy:** + - `useMemo` / `computed` / `derived` with correct dependency arrays + - Move derivation to selectors, loaders, or server-side + - Virtualize long lists (windowing) + - Stabilize callbacks and object props only when child renders are affected +- **Correctness checks:** Dependency arrays must include every semantic input; memoization must not hide mutations of mutable inputs + +#### Pairwise Comparison +Comparing every element with every other element using double-nested iteration. +- **Detection:** Two nested loops iterating the same or similar collections, comparing pairs +- **Impact:** O(n^2) for pair matching, overlap detection, conflict checking, nearest-neighbor +- **Remedy:** + - Sort + two-pointer for pair/range matching + - Sweep-line for interval overlaps + - Spatial hashing or grid bucketing for proximity + - Union-find for connectivity +- **Correctness checks:** Order stability; tie-breaking in equality cases + +#### Unnecessary Recompute (Missing Memoization) +Same pure computation repeated with same inputs without caching. +- **Detection:** Identical function calls with same arguments in hot paths; repeated expensive transforms; recursive calls without memoization +- **Impact:** Linear/polynomial wasted work; especially bad with recursive Fibonacci-style patterns (O(2^n) → O(n) with memo) +- **Remedy:** Add memoization/caching with proper invalidation; use `lru_cache`/`memoize`/`useMemo` as appropriate + +#### Wrong Data Structure +Using a suboptimal data structure for the access pattern. +- **Detection:** + - Array/List used for frequent membership tests → should be Set + - Array/List used for key-value lookups → should be Map/Object + - Array used as queue with `shift()`/`pop(0)` (O(n) per dequeue) → should use proper Queue + - Sorted insertion into array (O(n) per insert) → should use Heap +- **Remedy:** Replace with the data structure whose complexity matches the access pattern: + - Set → O(1) has/add/delete + - Map → O(1) get/set + - Heap → O(log n) push/pop for priority + - Queue/Deque → O(1) enqueue/dequeue + +#### What NOT to Flag +- **Cold paths:** Complexity that only runs on startup, config loading, or tiny N (< 100) is rarely worth fixing +- **Intentional tradeoffs:** Clear, readable O(n) code where O(n log n) would add complexity with no measurable gain +- **Already optimized:** Map/Set already in use; batch loading already implemented; memoization already present + +### Design Principle Violations + +#### SOLID Violations Checklist +- **S (SRP):** Class/module has multiple reasons to change → God Object smell +- **O (OCP):** switch/if-else chains on type codes → Strategy/Polymorphism needed +- **L (LSP):** Subclass changes behavior of base class unexpectedly → Check pre/post conditions +- **I (ISP):** Fat interfaces with methods clients don't use → Split interfaces +- **D (DIP):** High-level modules depending on low-level details → Introduce abstractions + +#### Other Principle Violations +- **DRY Violation:** Same knowledge repeated in multiple places +- **KISS Violation:** Over-engineered solutions; premature abstractions +- **YAGNI Violation:** Code for hypothetical future requirements; unused abstractions + +#### Object-Orientation Abusers (from Fowler / refactoring.guru) + +##### Switch Statements (Type-Code Conditionals) +Repeated `switch`/if-else chains that branch on a type code or enum: +- The same conditional structure duplicated in several places +- Adding a new type forces editing every switch (OCP violation) +- **Remedy:** Replace conditional with polymorphism (Strategy/State), or Replace Type Code with Subclasses + +##### Refused Bequest +A subclass inherits methods/fields it doesn't need: +- Overrides inherited methods to throw, no-op, or do something unrelated +- Signals the inheritance relationship is wrong +- **Remedy:** Push down unused members, or replace inheritance with delegation + +##### Alternative Classes with Different Interfaces +Two classes perform the same role but expose differently-named methods: +- `sort()` vs `arrange()`, `getUser()` vs `fetchUser()` for interchangeable classes +- **Remedy:** Unify the interface (rename methods, extract a common superclass/interface) + +##### Incomplete Library Class +A third-party/library class lacks methods you need and can't be modified: +- Scattered helper functions or copy-paste wrappers around the library +- **Remedy:** Introduce a Foreign Method or wrap it in an adapter/local extension class + +> Other refactoring.guru smells are documented in their thematic sections above: +> **Divergent Change**, **Data Class**, **Lazy Class**, **Speculative Generality**, +> **Temporary Field**, **Parallel Inheritance Hierarchies**, **Inappropriate Intimacy**, +> **Message Chains**, and **Middle Man**. + +--- + +## Edge Cases & Fallback + +| Scenario | Handling | +|----------|----------| +| User doesn't specify scope | Default to recent changes (`git diff`) for repos > 200 files, full analysis otherwise | +| Project has no clear architecture | Report "Big Ball of Mud" with evidence, recommend incremental refactoring | +| Empty/monorepo project | Report that architecture analysis requires code; ask user to specify module | +| Language not supported | Report general structural observations; note language-specific checks are limited | +| Report file path conflicts | Append `-2`, `-3`, etc. to filename | +| User wants a quick check | Run only Critical-level scans, skip Code and Naming categories | +| User wants only one category | Focus analysis on that category, skip others | + +--- + +## Report Output Example + +``` +🔍 Architecture Smell Analysis Complete + +Project: goal-workflow +Style: Modular Monolith (with some layering violations) +Files Analyzed: 47 +Health: 🟡 Fair + +Critical: 3 | Warnings: 6 | Suggestions: 9 + +🔴 Critical Issues: + 1. Anemic Domain Model — `models/` classes have only getters/setters, + all logic in `services/`. Violates DDD Rich Domain Model principle. + 2. N+1 Query Pattern — `services/order.ts:142` fetches user per order in loop; + should batch-load users by IDs (O(n*m) → O(n+m)). + 3. Static Cling — `util/ApiClient.ts` uses all static methods, + making consumer code untestable. + +🟡 Warnings: + 1. God Object — `services/workflow.ts` at 847 lines handles too many concerns + 2. Nested Loop O(n^2) — `analytics.ts:89` pairwise comparison of events; + sort+two-pointer would be O(n log n) + 3. Leaky Abstraction — `repositories/user.ts` exposes MongoDB query syntax + 4. Duplicated Code — validation logic duplicated across 4 controllers + 5. Circular Dependency — `auth` ↔ `user` modules depend on each other + 6. Magic Numbers — ~23 hardcoded values without named constants + +Full report: tasks/smell-report-2026-05-27-1530.md +``` \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/smell/test-prompts.json b/pigo/internal/builtinskills/skills/smell/test-prompts.json new file mode 100644 index 0000000..236dcb0 --- /dev/null +++ b/pigo/internal/builtinskills/skills/smell/test-prompts.json @@ -0,0 +1,50 @@ +[ + { + "prompt": "/smell", + "description": "Basic invocation — should trigger full project analysis" + }, + { + "prompt": "find code smells in this project", + "description": "Natural language trigger — full code smell analysis" + }, + { + "prompt": "detect architecture anti-patterns in the src/ directory", + "description": "Scoped analysis to a specific directory" + }, + { + "prompt": "分析一下这个项目的架构坏味道", + "description": "Chinese trigger — architecture smell analysis" + }, + { + "prompt": "找出反模式", + "description": "Chinese trigger — anti-pattern detection" + }, + { + "prompt": "Does this codebase have any God Objects or Big Ball of Mud?", + "description": "Specific anti-pattern query" + }, + { + "prompt": "run a quick architecture health check", + "description": "Quick scan — should use critical-only mode" + }, + { + "prompt": "check the recent changes for code smells", + "description": "Git diff scoped analysis" + }, + { + "prompt": "analyze code complexity and find algorithmic hotspots", + "description": "Complexity-focused analysis — detect N+1, nested loops, etc." + }, + { + "prompt": "find N+1 queries and O(n^2) patterns", + "description": "Specific complexity anti-pattern query" + }, + { + "prompt": "复杂度分析", + "description": "Chinese trigger — complexity analysis" + }, + { + "prompt": "Are there any performance bottlenecks or inefficient algorithms?", + "description": "Performance/complexity smell query" + } +] \ No newline at end of file diff --git a/pigo/internal/builtinskills/skills/to-design/SKILL.md b/pigo/internal/builtinskills/skills/to-design/SKILL.md new file mode 100644 index 0000000..9887681 --- /dev/null +++ b/pigo/internal/builtinskills/skills/to-design/SKILL.md @@ -0,0 +1,281 @@ +--- +name: to-design +description: "Generate a design document (design proposal) from a PRD, in the style of Go's official design proposals — Abstract / Background / Design / Rationale / Compatibility / Implementation, heavy on the 'why' and tradeoffs. Triggers on: to-design, prd-to-design, prd转设计文档, 生成设计文档, 写设计文档, design doc, design proposal, 设计提案, 技术设计文档." +user-invocable: true +--- + +# to-design — PRD to Design Document + +Turn a PRD (or a rough idea) into a **design document** written in the style of Go's official design proposals: plain language, concrete examples, and—above all—an honest account of *why this approach and not the alternatives*. + +This is **not** the same as `prd-to-spec`. A SPEC is an implementation contract (tables, endpoints, schemas) for an engineer to build against. A design document is a **decision artifact**: it argues for an approach, surfaces the tradeoffs, and lets a team agree on the same facts before anyone writes code. When the question is "*how should we build this and why*", produce a design doc; when the question is "*give me the exact contract to implement*", produce a SPEC. + +> 设计哲学源自对 5 篇 Go 官方 proposal(泛型 / 错误包装 / loopvar / slog / try)的分析。核心信念:**文档的价值不取决于方案是否通过,而取决于它是否让讨论建立在同一套事实和取舍之上。** + +--- + +## When to Use + +- A PRD exists and you need to decide *how* to build it before committing to implementation +- The approach has real tradeoffs and you want them documented and debated +- The change is risky, breaking, or hard to reverse (a design doc forces the compatibility conversation early) +- Multiple people need to agree on a direction before work fans out +- You want a durable record of "why we chose X and rejected Y" — even if the proposal is later rejected + +If the team just needs the concrete contract to code against, use `/prd-to-spec` instead (or run `to-design` first, then `prd-to-spec`). + +--- + +## The Job + +1. **Locate input** — find or receive the PRD (or idea) +2. **Analyze context (optional)** — scan the codebase for existing patterns, constraints, and prior art +3. **Surface the decisions** — identify the real design forks and ask clarifying questions (max 3-5) +4. **Generate the design doc** — following the structure and writing style below +5. **Review** — present for feedback, especially on the Rationale and Compatibility sections +6. **Save** — write to the agreed location + +--- + +## Step 1: Locate Input + +``` +Provide the PRD (or idea) to design from: + +A. File path (e.g., tasks/prd-priority-system.md) +B. GitHub Issue URL +C. Paste content directly +D. Just describe the idea — I'll design from the conversation +``` + +A design doc can start from a half-formed idea, not only a polished PRD. If the input is thin, lean harder on Step 3. + +--- + +## Step 2: Analyze Context (Optional) + +Skip for greenfield. Otherwise scan to ground the design in reality: + +- **Existing patterns** the design should match (naming, error handling, module boundaries) +- **Prior art** — has something similar been tried or rejected here before? +- **Constraints** — compatibility promises, public APIs, data the design can't break +- **Real pain** — find the actual buggy/awkward code the design fixes, so Background can quote it + +The most persuasive Background sections quote **real code from the user's own repo**, not hypotheticals. + +--- + +## Step 3: Surface the Decisions + +A design doc lives or dies on its Rationale. Before writing, find the **real forks in the road** — the points where a competent engineer could reasonably go two ways — and resolve them. + +Ask only about genuine forks: + +``` +Design decisions to settle before I write the doc: + +1. Where does this logic live? + A. Extend the existing X + B. New standalone component Y + C. Let me recommend based on the codebase + +2. Is this a breaking change for existing callers? + A. Yes — needs a migration path + B. No — purely additive + C. Unsure — I'll analyze and flag it + +3. What's the one promise this design must keep? (e.g. backward compatibility, + latency budget, no new dependencies) +``` + +For every fork, also note the **option you are NOT choosing** — that becomes the Rationale. + +--- + +## Step 4: Design Document Structure + +This is the standard skeleton distilled from the 5 Go proposals. Keep section names; drop sections that genuinely don't apply (and say why if the omission is notable). + +```markdown +Title: <一句话说清"做什么" —— 标题就是结论,不是名词短语> +Author(s): <作者> +Last updated: +Discussion at # 让文档不孤立,永远附讨论入口 +Status: Draft | Under review | Accepted | Rejected + +## Abstract / 摘要 + +一段话讲完全文:做什么、大致怎么做、以及**最重要的那个承诺**(如"向后兼容""不引入新依赖")。 +读者读完这一段就该知道全貌。把隐含的核心约束埋在这里。 + +## Background / 背景与动机 + +用**具体、可感的例子**说明"痛在哪",而不是抽象地说"现状不好"。 +- 能贴一段真实的 bug 代码 / 别扭的调用,就贴。先让读者"疼"起来。 +- 量化痛点(出现频率、踩坑次数、损失),不要用形容词堆砌。 +- 一句话给问题定性。 + +## Design / Proposal / 设计 + +文档主体。遵循三条: +- **从简单到复杂,渐进式教学**:从最小例子起步,复杂场景留到读者有直觉之后。 +- **声明 + 示例 + 边界**三件套:每个 API/接口先给声明,再给用法片段,再划清适用边界。 +- **改造前 vs 改造后对照**:能并排展示收益的,就并排展示。 +能用一段可运行代码说清的,绝不用一段文字描述。 + +## Rationale / 理由与取舍 + +> Rationale = "为什么是这个方案,而不是别的"的论证。这是区分好文档和平庸文档的关键章节。 + +- 解释关键决策的动机。 +- **主动列出被放弃的备选方案 + 放弃原因**("我们没选 X,因为 Y")。这比单方面论证你选的方案更可信,也避免后人重复讨论。 +- 回应可预见的质疑。 + +## Compatibility / 兼容性 + +凡涉及破坏性变更,必须正面回应。 +- 是不是破坏性变更?**开门见山承认**。 +- 代价是什么(性能、行为变化、迁移成本)?**诚实列出**,不藏着。 +- 渐进迁移路径(按模块/按文件 opt-in、灰度、特性开关)。 +- 有先例佐证更好("某系统做过类似变更,结果平淡无奇")。 + +## Implementation / Transition / 实现与过渡 + +- 如何落地、分几步、配套什么工具。 +- **用数据和工具支撑"可落地"**:实测失败率、灰度结果、自动化迁移工具,比任何"我们认为风险可控"都管用。 +- 兼容老版本的过渡方案(如独立发布的兼容库)。 + +## Appendix / 附录(可选) + +把会打断主线的细节后置:完整 API、端到端示例、FAQ。 +FAQ 专门回应高频质疑("为什么叫这个名字""为什么不用某语言的做法""和 X 有何不同")。 +``` + +--- + +## Writing Style (照搬 Go 文档的文风) + +Structure is the skeleton; style is the muscle. Enforce these — they're what make the doc readable. + +### Voice / 主语 +- **决策用 "我们 / We"** — 把设计说成一群人可负责的选择,不是客观真理。("We propose…", "我们决定移除…") +- **行为用代码本身当主语** — "this code has a bug" / "这段代码会…",让注意力落在程序上。 +- **说理对读者用 "你 / you"** — 像面对面解释。 +- **禁止无主语的被动腔** — 不写"据建议应当…""It is suggested that…"这类推卸责任的句式。 + +### Sentences / 句子 +- **判断用短句,论证用长句**。先用一个极短的句子拍板("这段代码有 bug。"),再用信息密集的长句铺开机制。 +- 长短交替制造节奏。不要通篇绕来绕去的长句。 + +### Paragraphs / 段落 +- **一段只讲一件事,观点放段首**(结论先行)。 +- **小标题写成一句完整的论点**,而不是名词短语。 + - 写 `老代码不受影响,编译结果与之前完全一致`,而不是 `兼容性`。 + - 读者光看标题就能读完整条论证链。 + +### Tone / 语气 +- **克制的诚实,甚至自嘲**。承认代价、承认自己也踩过坑,比形容词更有说服力。 +- **强调要省着用**。全文只在最关键处加粗/斜体一次,反而最醒目。 + +--- + +## Step 5: Review & Iteration + +Present the doc and steer feedback to the sections that matter most: + +``` +设计文档已生成。重点请看这几处: + +- Rationale:被放弃的方案和理由是否站得住?有没有遗漏的备选项? +- Compatibility:破坏性和代价是否如实说清?迁移路径可行吗? +- Background:痛点是否用具体例子讲清,而不是形容词? +- 文风:标题是否是"结论"而非名词?有没有无主语的被动腔? + +回复 OK 保存,或给出修改意见。 +``` + +--- + +## Step 6: Save + +``` +设计文档保存到哪里? + +A. tasks/design-[feature-name].md(紧挨 PRD,推荐) +B. docs/design/[feature-name].md +C. 自定义路径:[指定] +``` + +--- + +## Mapping: PRD → Design Doc + +| PRD 部分 | Design Doc 部分 | 转化方式 | +|----------|-----------------|----------| +| Problem / 背景 | Background | 找到真实的痛点代码/场景,量化它 | +| Goals / 目标 | Abstract + Background | 提炼成"最重要的承诺"埋进摘要 | +| User Stories / 需求 | Design | 转成渐进式的设计示例 | +| Technical Considerations | Design + Rationale | 约束 → 设计决策 + 取舍论证 | +| Non-Goals | Rationale | 写成"我们没做 X,因为 Y" | +| Risks / 风险 | Compatibility + Implementation | 风险 → 兼容性代价 + 迁移/灰度方案 | +| 隐含的备选方案 | Rationale | 显式列出并解释为何不选 | + +--- + +## Quality Criteria + +A good design doc should pass these checks: + +- [ ] 标题是一句"做什么"的结论,不是名词短语,且附了讨论链接 +- [ ] 摘要里埋了最重要的承诺/约束 +- [ ] Background 用了**具体例子或真实代码**讲痛点,而非形容词 +- [ ] Design 遵循"声明 + 示例 + 边界",并有渐进式教学 +- [ ] **Rationale 主动列出了至少一个被放弃的方案及原因**(最关键的检查项) +- [ ] 凡破坏性变更,Compatibility 都正面承认并列出代价 +- [ ] Implementation 用数据/工具支撑"可落地",而非空喊"风险可控" +- [ ] 文风:决策用"我们"、行为用代码、无无主语被动腔;长短句交替;小标题是论点句 +- [ ] 没有 "TBD / TODO"——要么解决,要么挪进 Open Questions + +--- + +## Edge Cases & Fallback + +| 场景 | 处理 | +|------|------| +| PRD 含糊不全 | 在 Step 3 多问,把缺失项写进 Open Questions / 假设 | +| 没有真实痛点代码可引 | 用最小可信的示例代码代替,并注明是构造的 | +| 没有备选方案可写 | 强迫思考"最朴素的做法是什么、为什么不够"——总有一个被否决的基线 | +| 不是破坏性变更 | Compatibility 一句话说明"纯增量、无破坏",不必硬凑 | +| 方案最终被否决 | 照样写好——记录"这条路为什么走不通"本身就是高价值产物,Status 标 Rejected | +| 特性太大 | 拆成多篇 design doc(按边界),互相链接 | +| 用户只要实现契约 | 提示改用 `/prd-to-spec`,或先 to-design 再 prd-to-spec | + +--- + +## Anti-Patterns to Avoid + +- **别只论证你选的方案。** 不写被放弃的备选项,文档就少了一半价值。 +- **别用形容词讲痛点。** "现状很糟"没有说服力;一段真实的 bug 代码才有。 +- **别藏代价。** 性能变慢、行为变化、迁移成本——都明说,再给迁移路径。 +- **别把标题写成名词。** "兼容性" → "老代码不受影响,编译结果完全一致"。 +- **别用无主语的被动腔。** 决策要有人负责,主语用"我们"。 +- **别写成 SPEC。** 设计文档讲"为什么这么选"和"取舍",不是字段级的实现契约。 +- **别因为方案可能被否就敷衍。** 文档质量与提案是否通过无关。 + +--- + +## Relationship to Other Skills + +``` +/prd → /to-design → /prd-to-spec → /goal → /review-it → /ship-it + │ │ │ │ + │ 需求(what) │ 决策与取舍 │ 实现契约(how) │ 编码 + │ │ (why/which) │ +``` + +- **/prd** 产出 PRD(本 skill 的输入) +- **/to-design** 产出设计文档:论证方案、暴露取舍、对齐认知(本 skill) +- **/prd-to-spec** 产出实现级 SPEC:字段、接口、schema 契约 +- **/code-to-spec** 从既有代码逆向出 SPEC(互补:正向 vs 逆向) + +> 写设计文档的终极目的不是"说服别人同意你",而是"让所有人在同一个事实和取舍基础上做决定"。 diff --git a/pigo/internal/builtinskills/skills/to-issues/SKILL.md b/pigo/internal/builtinskills/skills/to-issues/SKILL.md new file mode 100644 index 0000000..647375f --- /dev/null +++ b/pigo/internal/builtinskills/skills/to-issues/SKILL.md @@ -0,0 +1,229 @@ +--- +name: to-issues +description: "Decompose a PRD and/or SPEC into implementable Issues and create them in your chosen platform (GitHub, Local, or Baidu iCafe). Use after /prd (and optionally /prd-to-spec) to turn requirements into actionable tickets. Triggers on: create issues, to-issues, 创建issue, 拆解issue, 生成卡片, 创建卡片, generate issues from PRD, issues from spec." +user-invocable: true +--- + +# to-issues — PRD/SPEC to Issues + +Decompose a PRD and/or technical SPEC into small, independent, implementable Issues, then create them in your chosen platform. Works standalone — you don't need to have run `/prd` first. + +--- + +## The Job + +1. **Locate input** — find a PRD or SPEC file (auto-detect or user-specified) +2. **Decompose into Issues** — break User Stories into implementable tickets +3. **Review with user** — present Issue list for approval and adjustment +4. **Choose platform** — GitHub / Local / Baidu iCafe +5. **Create Issues** — create all tickets and print summary + +--- + +## Step 1: Locate Input + +Find the input document: + +``` +What should I base the Issues on? + +A. Auto-detect: scan tasks/ for recent PRDs and SPECs +B. Specific PRD file (e.g., tasks/prd-priority-system.md) +C. Specific SPEC file (e.g., tasks/spec-priority-system.md) +D. Both PRD and SPEC (best: PRD for requirements, SPEC for technical contracts) +E. Paste requirements directly +``` + +If auto-detecting, list available files and let the user choose. + +If both PRD and SPEC are available, use the SPEC's Section 10.2 (Issue Mapping) as the primary guide, supplemented by PRD's User Stories. If only PRD is available, generate Issues directly from User Stories. + +--- + +## Step 2: Decompose into Issues + +Based on the input document(s), generate a list of Issues. Follow these rules: + +- **One Issue per User Story** — each US-XXX becomes at least one Issue +- **Split large stories** — if a US has 5+ acceptance criteria or spans frontend + backend, split into 2-3 smaller Issues with clear dependencies +- **Merge tiny stories** — if a US has only 1-2 trivial criteria, merge it with a related US into a single Issue +- **Each Issue must be independently implementable** — a single agent session should be able to complete it +- **Number Issues sequentially** starting from 1 +- **If SPEC is available** — enrich Issues with SPEC references (API endpoints, data model sections, error handling contracts) + +**Issue format:** + +``` +Issue #N: [Title] +--- +Description: [From US description, with context] +Acceptance Criteria: +- [ ] [From US acceptance criteria] +- [ ] ... +Dependencies: [None / Issue #X] +Type: [backend / frontend / fullstack / ui / infra] +Priority: [high / medium / low] +SPEC Reference: [Section X.Y — only if SPEC available] +``` + +**Present the Issue list for review:** + +``` +📋 Generated N Issues from [PRD/SPEC]: + +#1: Add priority field to database (backend, high) +#2: Display priority indicator on task cards (frontend, high) — depends on #1 +#3: Add priority selector to task edit (frontend, medium) — depends on #1 +#4: Filter tasks by priority (frontend, medium) — depends on #1, #2 + +Please review. You can: +- Remove issues: "remove #3" +- Merge issues: "merge #2 and #3" +- Add issues: "add an issue for sorting by priority" +- Adjust: "change #2 priority to high" +- Confirm: reply OK to proceed +``` + +Wait for user confirmation before creating any Issues. + +--- + +## Step 3: Choose Creation Mode + +After user confirms the Issue list, ask: + +``` +Choose where to create these Issues: + +A. GitHub (via gh CLI) +B. Local (save as .md files) +C. Baidu iCafe (via icafe-cli) + +Your choice: +``` + +--- + +## Step 4: Mode-Specific Creation + +### Mode A: GitHub + +**Prerequisites:** `gh` CLI installed and authenticated. + +**Actions:** +1. For each Issue, run: + ```bash + gh issue create --title "[Title]" --body "[Description + Acceptance Criteria]" --label "[type]" --label "priority: [priority]" + ``` +2. If labels don't exist, create them first or skip the `--label` flag +3. Report created Issue numbers and URLs + +### Mode B: Local + +**Ask user:** +``` +Where should I save the Issue files? (default: .autoresearch/issues) +``` + +**Actions:** +1. If the specified folder does not exist, create it with `mkdir -p` +2. For each Issue #N, save a file named `issue-NNN-[slug].md` (zero-padded to 3 digits): + ```markdown + # [Title] + + ## Description + [Description from Issue] + + ## Acceptance Criteria + - [ ] [criterion 1] + - [ ] [criterion 2] + + ## Dependencies + [None / Issue #X] + + ## Type + [backend / frontend / fullstack / ui / infra] + + ## Priority + [high / medium / low] + ``` +3. Report created file paths + +### Mode C: Baidu iCafe + +**Ask user:** +``` +Please provide the iCafe space prefix code (--space): +``` + +Optionally ask: +``` +Target branch for iCode CR? (default: master) +``` + +**Prerequisites:** `icafe-cli` installed and logged in. + +**Actions:** +1. For each Issue, run: + ```bash + icafe-cli card create --space [SPACE] --title "[Title]" --description "[Description + Acceptance Criteria]" --cardtype "[Task/Bug/Story]" + ``` + - Map Issue `type` to iCafe card type: `bug` → `Bug`, `ui`/`frontend` → `Story`, others → `Task` + - Map `priority`: high → `高`, medium → `中`, low → `低` +2. If iCafe card creation fails for an Issue, log the error and continue with remaining Issues +3. Report created card sequence numbers + +--- + +## Step 5: Summary Report + +After all Issues are created, print a summary: + +``` +✅ Issue creation complete! + +Source: [PRD/SEC path] +Mode: [GitHub / Local / Baidu iCafe] +Issues created: N + +# | Title | Identifier +---|------------------------------------------|------------ +1 | Add priority field to database | #42 (GitHub) / issue-001-*.md (Local) / #22210 (iCafe) +2 | Display priority indicator | #43 / issue-002-*.md / #22211 +3 | Add priority selector | #44 / issue-003-*.md / #22212 +4 | Filter tasks by priority | #45 / issue-004-*.md / #22213 + +💡 Tip: Now implement each Issue with /goal: + /goal 42 # GitHub mode + /goal issue-001-*.md # Local mode +``` + +--- + +## Edge Cases & Fallback + +| Scenario | Handling | +|----------|----------| +| No PRD/SPEC found in tasks/ | Ask user to provide file path or paste requirements | +| PRD has no User Stories | Derive Issues from Functional Requirements instead | +| SPEC has Issue Mapping (Section 10.2) | Use it as primary source, cross-reference with PRD | +| `gh` CLI not authenticated for GitHub mode | Show error, suggest `gh auth login`, offer to switch to Local mode | +| `icafe-cli` / `icode-cli` not installed for Baidu mode | Show error, suggest installation, offer to switch to Local mode | +| Issue folder does not exist for Local mode | Auto-create the folder | +| User declines Issue creation | Print the Issue list as a text summary, let user create manually later | + +--- + +## Relationship to Other Skills + +``` +/prd → /prd-to-spec (optional) → /to-issues → /goal → /review-it → /ship-it + │ │ │ │ + │ Requirements │ Technical design │ Tickets │ Implementation + │ (what) │ (how) │ (units) │ (code) +``` + +- **/prd** — produces the PRD (input to this skill) +- **/prd-to-spec** — produces the SPEC (optional, enriches Issues with technical detail) +- **/to-issues** — produces the Issues (this skill) +- **/goal** — implements Issues one by one diff --git a/pigo/internal/builtinskills/skills/weather/SKILL.md b/pigo/internal/builtinskills/skills/weather/SKILL.md new file mode 100644 index 0000000..2146580 --- /dev/null +++ b/pigo/internal/builtinskills/skills/weather/SKILL.md @@ -0,0 +1,49 @@ +--- +name: weather +description: Get current weather and forecasts (no API key required). +homepage: https://wttr.in/:help +metadata: {"clawdbot":{"emoji":"🌤️","requires":{"bins":["curl"]}}} +--- + +# Weather + +Two free services, no API keys needed. + +## wttr.in (primary) + +Quick one-liner: +```bash +curl -s "wttr.in/London?format=3" +# Output: London: ⛅️ +8°C +``` + +Compact format: +```bash +curl -s "wttr.in/London?format=%l:+%c+%t+%h+%w" +# Output: London: ⛅️ +8°C 71% ↙5km/h +``` + +Full forecast: +```bash +curl -s "wttr.in/London?T" +``` + +Format codes: `%c` condition · `%t` temp · `%h` humidity · `%w` wind · `%l` location · `%m` moon + +Tips: +- URL-encode spaces: `wttr.in/New+York` +- Airport codes: `wttr.in/JFK` +- Units: `?m` (metric) `?u` (USCS) +- Today only: `?1` · Current only: `?0` +- PNG: `curl -s "wttr.in/Berlin.png" -o /tmp/weather.png` + +## Open-Meteo (fallback, JSON) + +Free, no key, good for programmatic use: +```bash +curl -s "https://api.open-meteo.com/v1/forecast?latitude=51.5&longitude=-0.12¤t_weather=true" +``` + +Find coordinates for a city, then query. Returns JSON with temp, windspeed, weathercode. + +Docs: https://open-meteo.com/en/docs diff --git a/pigo/internal/builtinskills/skills/weather/_meta.json b/pigo/internal/builtinskills/skills/weather/_meta.json new file mode 100644 index 0000000..4556002 --- /dev/null +++ b/pigo/internal/builtinskills/skills/weather/_meta.json @@ -0,0 +1,6 @@ +{ + "ownerId": "kn70pywhg0fyz996kpa8xj89s57yhv26", + "slug": "weather", + "version": "1.0.0", + "publishedAt": 1767545394459 +} \ No newline at end of file diff --git a/pigo/internal/cli/btw/btw.go b/pigo/internal/cli/btw/btw.go new file mode 100644 index 0000000..824d849 --- /dev/null +++ b/pigo/internal/cli/btw/btw.go @@ -0,0 +1,302 @@ +// This file implements the /btw command (mirrors Claude Code's /btw and the pi +// agent extension @narumitw/pi-btw): a throwaway "side thread" for asking the +// model a quick side question that must NOT pollute the main conversation. +// +// /btw is intercepted in the REPL loop rather than routed through a slash Action +// closure because it must run an agent stream and read the live main context — +// none of which a pure string→string Action can do, exactly like /compact and +// /goal. It reaches the session's collaborators and mutable state through the +// cli.Host contract and reads follow-up lines through cli.Editor, so it need not +// import the concrete replDeps aggregate that assembles them. +// +// Isolation contract (the whole point of the feature): a side thread runs on a +// COPY of the main conversation as background, and its question/answer are only +// ever appended to that copy — never to host.AgentCtx().Messages. Nothing is +// persisted: no store.Save, no change to the persisted cursor / current leaf / +// header timestamp. Closing the side thread, switching sessions or restarting +// pigo discards everything. +// +// Scope: /btw is intercepted in the REPL loop and runs a side question against a +// copy of the main context (#279); it supports multi-turn follow-ups in the same +// ephemeral thread (#280), bare-/btw reopen of the most recent side thread this +// process (#281), and an optional model/thinking override config (#282, see +// btw_config.go) that affects only the side thread. +package btw + +import ( + "context" + "errors" + "fmt" + "io" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/trust" +) + +// btwHeader is the fixed banner shown when entering a side thread, so the user +// always knows the current input is a throwaway side question, not the main +// conversation (mirrors pi-btw's "btw · side thread" header). +const btwHeader = "btw · side thread" + +// BtwHeader exposes the side-thread banner text so callers (and tests) can +// recognize it in output. +const BtwHeader = btwHeader + +// btwPrompt is the input prompt shown for follow-up questions inside a side +// thread, distinguishing it from the main "pigo(model)>" prompt. +const btwPrompt = "btw> " + +// RunBtw handles a /btw invocation. With an argument it starts a fresh side +// thread, asks that question, then enters a follow-up loop so the user can keep +// asking in the same ephemeral thread. Bare "/btw" reopens the most recent side +// thread from this process — replaying its Q&A history — and drops back into the +// follow-up loop; if none exists yet it guides the user to supply a question +// (US-004, #281). setCancel publishes the active run's cancel func so the REPL's +// SIGINT handler can interrupt the side run, reusing the same plumbing as a +// normal turn. +// +// The main context is never mutated: RunBtw builds a private side AgentContext +// seeded with a copy of the main messages, runs every turn against that copy, +// and returns without touching host.AgentCtx() or persisting anything. The side +// thread is retained in-process (host.LastBtw()) so a later bare /btw can reopen +// it, but it is never written to disk — restarting pigo discards it. +func RunBtw(setCancel func(context.CancelFunc), out io.Writer, host cli.Host, editor cli.Editor, line string) { + question := strings.TrimSpace(strings.TrimPrefix(line, "/btw")) + // Resolve the side thread's model/thinking once per invocation from the + // session defaults overlaid with btw.json (#282). Re-read each call so an + // edit takes effect next time with no restart. + settings := ResolveBtwSettings(out, host) + if question == "" { + // Bare /btw: reopen the most recent side thread if one exists this process, + // replaying its history; otherwise guide the user to supply a question. + if host.LastBtw() == nil { + fmt.Fprintln(out, "usage: /btw — ask a quick side question without touching the main conversation") + return + } + printBtwHeader(out) + replaySideHistory(out, host.LastBtw(), host.LastBtwBase()) + if editor != nil { + btwFollowUpLoop(setCancel, out, host, editor, host.LastBtw(), settings) + } + return + } + + side := NewSideContext(host.AgentCtx()) + // Remember this thread so a later bare /btw can reopen it. LastBtwBase marks + // where the copied background ends and the side Q&A begins, so a reopen only + // replays the side turns, not the whole main transcript. + host.SetLastBtw(side) + host.SetLastBtwBase(len(side.Messages)) + printBtwHeader(out) + AskSide(setCancel, out, host, side, settings, question) + // Follow-up loop: keep answering in the same ephemeral thread until the user + // exits. A nil editor (direct test callers that only ask one question) skips + // the loop entirely, so a single /btw asks exactly one question and returns. + if editor != nil { + btwFollowUpLoop(setCancel, out, host, editor, side, settings) + } +} + +// replaySideHistory prints the side thread's own Q&A (everything after the +// copied main-conversation background at index base) when a bare /btw reopens a +// prior thread, so the user can browse earlier answers before continuing. Only +// user questions and assistant text are shown; tool activity is omitted to keep +// the recap compact. +func replaySideHistory(out io.Writer, side *agentcore.AgentContext, base int) { + if base > len(side.Messages) { + base = len(side.Messages) + } + for _, msg := range side.Messages[base:] { + switch m := msg.(type) { + case agentcore.UserMessage: + fmt.Fprintf(out, "%s %s\n", ui.Colorize(ui.Enabled(), ui.Dim, "you:"), agentcore.ContentToText(m.Content)) + case agentcore.AssistantMessage: + if text := agentcore.ContentToText(m.Content); text != "" { + rendered := ui.RenderMarkdown(text) + fmt.Fprint(out, rendered) + if !strings.HasSuffix(rendered, "\n") { + fmt.Fprintln(out) + } + } + } + } +} + +// btwFollowUpLoop reads follow-up questions and answers them in the same side +// context, so each answer sees the prior side Q&A (FR-4). It exits on /exit, +// /quit, EOF, or an idle Ctrl+C (errLineInterrupted) — the same exit affordances +// as the main REPL, but confined to the side thread (FR-5). A blank line is +// ignored (stays in the thread). Nothing here touches the main context. +func btwFollowUpLoop(setCancel func(context.CancelFunc), out io.Writer, host cli.Host, editor cli.Editor, side *agentcore.AgentContext, settings BtwRunSettings) { + for { + raw, err := editor.ReadLine(btwPrompt) + if errors.Is(err, cli.ErrLineInterrupted) { + // Idle Ctrl+C at the side prompt leaves the thread (a Ctrl+C during a + // run is handled inside askSide via the SIGINT cancel plumbing). + fmt.Fprintln(out, "left side thread") + return + } + q := strings.TrimSpace(raw) + if err != nil && q == "" { + // EOF or read error with no partial line: leave the thread. + fmt.Fprintln(out, "left side thread") + return + } + if q == "/exit" || q == "/quit" { + fmt.Fprintln(out, "left side thread") + return + } + if q == "" { + continue + } + AskSide(setCancel, out, host, side, settings, q) + } +} + +// NewSideContext builds the side thread's private AgentContext. Its Messages are +// a fresh slice seeded with a shallow COPY of the main messages (the elements +// are immutable value/interface messages, so a copied slice header is enough to +// guarantee appends to the side thread never reach the main context's Messages). +// The system prompt and tools are shared by value; only Messages diverges. +func NewSideContext(main *agentcore.AgentContext) *agentcore.AgentContext { + msgs := make(agentcore.MessageList, len(main.Messages)) + copy(msgs, main.Messages) + return &agentcore.AgentContext{ + SystemPrompt: main.SystemPrompt, + Messages: msgs, + Tools: main.Tools, + } +} + +// printBtwHeader prints the side-thread banner. +func printBtwHeader(out io.Writer) { + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, btwHeader)) +} + +// AskSide appends the question to the side context and streams one answer, +// mirroring streamRun's rendering but targeting the side context so nothing is +// written back to the main conversation or to disk. It reuses the REPL's SIGINT +// cancel plumbing via setCancel. The model/provider/thinking come from settings +// (session defaults overlaid with btw.json, #282), never from host.Live(), so a +// /btw override cannot leak into the main session. +func AskSide(setCancel func(context.CancelFunc), out io.Writer, host cli.Host, side *agentcore.AgentContext, settings BtwRunSettings, question string) { + content, err := ui.BuildUserContent(question) + if err != nil { + fmt.Fprintf(out, "pigo: %v\n", err) + return + } + side.Messages = append(side.Messages, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: content, + }) + + runCtx, cancel := context.WithCancel(context.Background()) + setCancel(cancel) + defer func() { + cancel() + setCancel(nil) + }() + + // Show a transient status while the model works (FR-9). It is printed on its + // own line; the streamed answer follows below it. + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, "Answering…")) + + cfg := runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: settings.Model, + Provider: settings.ProviderName, + ThinkingLevel: settings.ThinkingLevel, + Stream: provider.StreamFnFromProvider(settings.Provider), + GetAPIKey: host.Creds().GetAPIKey, + ContextWindow: host.Live().ContextWindow, + Compaction: compaction.DefaultCompactionSettings, + }, + Batch: agenttool.BatchConfig{ + ToolExecutorConfig: agenttool.ToolExecutorConfig{ + Registry: host.Registry(), + BeforeToolCall: trust.BeforeToolCall(host.Trust(), host.Cwd(), host.Input(), out, host.ConfirmMu()), + }, + }, + Reminders: host.Reminders(), + } + // Wire the per-turn hook seams onto the side run's cfg; nil dispatcher is a + // no-op (FR-18). + if d := host.Dispatcher(); d != nil { + run.InstallSeams(&cfg, d, host.HookDeps()) + } + stream := runtime.StartRun(runCtx, side, cfg) + drainSideStream(runCtx, out, host, stream) +} + +// chainBtwEvent returns the OnEvent observer for a /btw side run: the plugin +// notifier, with the SessionEnd/PreCompact hook notifier chained after it when +// hooks are configured, mirroring the REPL's OnEvent composition. +func chainBtwEvent(host cli.Host) func(agentcore.AgentEvent) { + notifier := host.NotifierHandle() + d := host.Dispatcher() + if d == nil { + return notifier + } + deps := host.HookDeps() + hookEvent := hooks.NewHookNotifier(d, deps.SessionID, deps.ProjectDir).Handle + if notifier == nil { + return hookEvent + } + return func(ev agentcore.AgentEvent) { + notifier(ev) + hookEvent(ev) + } +} + +// drainSideStream prints the streamed assistant text and tool activity of a side +// run, mirroring streamRun/drainGoalStream. It blocks until the run ends. Unlike +// the main loop it persists nothing. +func drainSideStream(ctx context.Context, out io.Writer, host cli.Host, stream *runtime.LoopEventStream) { + var reply strings.Builder + flushReply := func() { + if reply.Len() == 0 { + return + } + rendered := ui.RenderMarkdown(reply.String()) + fmt.Fprint(out, rendered) + if !strings.HasSuffix(rendered, "\n") { + fmt.Fprintln(out) + } + reply.Reset() + } + _, err := runtime.DrainStream(ctx, stream, runtime.StreamHandler{ + OnEvent: chainBtwEvent(host), + OnText: func(delta string) { + reply.WriteString(delta) + }, + OnTurnEnd: func(msg agentcore.AssistantMessage, results []agentcore.ToolResultMessage) { + flushReply() + for _, c := range msg.ToolCalls() { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(ui.Enabled(), ui.Green, "→ tool:"), ui.ToolCallLabel(c)) + } + for _, tr := range results { + ui.RenderToolResult(out, tr) + } + }, + }) + flushReply() + if err != nil { + if ctx.Err() != nil { + // A Ctrl+C during the run cancels just this answer; the follow-up loop + // then returns to the btw prompt so the user can ask again or exit with + // another Ctrl+C (FR-5). + fmt.Fprintln(out, "^C interrupted — answer cancelled") + } else { + fmt.Fprintf(out, "error: %v\n", err) + } + } +} diff --git a/pigo/internal/cli/btw/btw_config.go b/pigo/internal/cli/btw/btw_config.go new file mode 100644 index 0000000..2e0751a --- /dev/null +++ b/pigo/internal/cli/btw/btw_config.go @@ -0,0 +1,140 @@ +// This file implements the /btw model/thinking override config (US-005, #282): +// an optional per-command config that lets a side thread use a different model +// and/or reasoning effort than the main session, without touching the main +// session's settings (mirrors pi-btw's pi-btw.json). +// +// The config lives at $PIGO_HOME/btw.json (or ~/.pigo/btw.json). It is read +// fresh on every /btw invocation, so editing it takes effect on the next call +// with no restart. A missing file, an empty object, or an absent field all mean +// "inherit the session default" silently — only a malformed file or an +// unusable model override produces a (non-fatal) warning. +package btw + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/provider" +) + +// btwConfig is the on-disk shape of ~/.pigo/btw.json. Both fields are optional; +// an absent field (nil / empty) inherits the session default. Pointers/empty +// strings distinguish "not set" from a real value so a partial file still falls +// back per-field. +type btwConfig struct { + Model string `json:"model,omitempty"` + ThinkingLevel string `json:"thinkingLevel,omitempty"` +} + +// BtwRunSettings is the resolved model/provider/thinking a side run uses. It is +// computed once per /btw invocation from the session defaults overlaid with +// btw.json, and passed down to AskSide so every turn of that invocation uses +// the same settings. +type BtwRunSettings struct { + Model string + ProviderName string + Provider provider.Provider + ThinkingLevel agentcore.ThinkingLevel +} + +// btwConfigPath returns the path to the /btw override config, or "" when the +// config directory cannot be resolved (then the config is treated as absent). +func btwConfigPath() string { + dir := run.ConfigDir() + if dir == "" { + return "" + } + return filepath.Join(dir, "btw.json") +} + +// loadBtwConfig reads and parses btw.json. A missing file returns a zero config +// with no error (inherit everything). A malformed file returns an error so the +// caller can warn and fall back. An empty object parses to a zero config. +func loadBtwConfig(path string) (btwConfig, error) { + if path == "" { + return btwConfig{}, nil + } + data, err := os.ReadFile(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return btwConfig{}, nil + } + return btwConfig{}, err + } + var cfg btwConfig + if err := json.Unmarshal(data, &cfg); err != nil { + return btwConfig{}, fmt.Errorf("parse %s: %w", path, err) + } + return cfg, nil +} + +// ResolveBtwSettings computes the model/provider/thinking a side run should use. +// It starts from the session defaults (host.Live()) and overlays btw.json: +// +// - No config / empty object / absent fields → inherit the session values. +// - thinkingLevel set → validate and override (invalid value warns, falls back). +// - model set → resolve its provider (reusing resolveProvider like /model); +// if the model cannot be resolved/authenticated, warn on one line and fall +// back to the session model+provider. +// +// A malformed config file warns once and inherits everything. Nothing here +// mutates the session live config, so the override is confined to the side +// thread (FR-8). +func ResolveBtwSettings(out io.Writer, host cli.Host) BtwRunSettings { + live := host.Live() + s := BtwRunSettings{ + Model: live.Model, + ProviderName: live.ProviderName, + Provider: live.Provider, + ThinkingLevel: live.ThinkingLevel, + } + + cfg, err := loadBtwConfig(btwConfigPath()) + if err != nil { + fmt.Fprintf(out, "%s\n", ui.Colorize(ui.Enabled(), ui.Dim, "btw: ignoring invalid btw.json: "+err.Error())) + return s + } + + if lvl := strings.TrimSpace(cfg.ThinkingLevel); lvl != "" { + if v, ok := validThinkingLevel(lvl); ok { + s.ThinkingLevel = v + } else { + fmt.Fprintf(out, "%s\n", ui.Colorize(ui.Enabled(), ui.Dim, fmt.Sprintf("btw: ignoring invalid thinkingLevel %q, using %q", lvl, s.ThinkingLevel))) + } + } + + if model := strings.TrimSpace(cfg.Model); model != "" && model != s.Model { + prov, providerName, perr := provider.ResolveProvider(model, live.BaseURL, live.Protocol, "", os.Getenv) + if perr != nil { + fmt.Fprintf(out, "%s\n", ui.Colorize(ui.Enabled(), ui.Dim, fmt.Sprintf("btw: cannot use model %q (%v), falling back to %q", model, perr, s.Model))) + } else { + s.Model = model + s.ProviderName = providerName + s.Provider = prov + } + } + + return s +} + +// validThinkingLevel reports whether s is one of the known reasoning-effort +// levels and returns the typed value. It mirrors the enum in agentcore so an +// invalid btw.json value can be rejected without importing the config layer. +func validThinkingLevel(s string) (agentcore.ThinkingLevel, bool) { + switch agentcore.ThinkingLevel(s) { + case agentcore.ThinkingOff, agentcore.ThinkingMinimal, agentcore.ThinkingLow, + agentcore.ThinkingMedium, agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax: + return agentcore.ThinkingLevel(s), true + default: + return "", false + } +} diff --git a/pigo/internal/cli/btw/btw_config_test.go b/pigo/internal/cli/btw/btw_config_test.go new file mode 100644 index 0000000..d3f98f6 --- /dev/null +++ b/pigo/internal/cli/btw/btw_config_test.go @@ -0,0 +1,146 @@ +package btw + +// Tests for the /btw model/thinking override config (#282, US-005): btw.json +// overlays the session defaults for the side thread only, is read fresh each +// call, and falls back silently on missing/empty/partial config. + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" +) + +// fakeHost satisfies cli.Host by embedding the interface (so every method is +// present) while overriding only Live(), the sole accessor ResolveBtwSettings +// reads. The embedded nil interface would panic if any other method were +// called, which these tests never do. +type fakeHost struct { + cli.Host + live *cli.LiveConfig +} + +func (f fakeHost) Live() *cli.LiveConfig { return f.live } + +// withBtwConfig points PIGO_HOME at a temp dir and writes btw.json with the +// given contents (or removes it when contents is ""), returning nothing — the +// temp dir is cleaned up by t.TempDir. It restores PIGO_HOME after the test. +func withBtwConfig(t *testing.T, contents string) { + t.Helper() + dir := t.TempDir() + t.Setenv("PIGO_HOME", dir) + if contents != "" { + if err := os.WriteFile(filepath.Join(dir, "btw.json"), []byte(contents), 0o644); err != nil { + t.Fatalf("write btw.json: %v", err) + } + } +} + +// TestBtwConfigAbsentInherits verifies that with no btw.json the side settings +// equal the session defaults. +func TestBtwConfigAbsentInherits(t *testing.T) { + withBtwConfig(t, "") // no file + live := &cli.LiveConfig{Model: "sess-model", ProviderName: "sess-prov", ThinkingLevel: agentcore.ThinkingMedium} + host := fakeHost{live: live} + + var warn bytes.Buffer + s := ResolveBtwSettings(&warn, host) + if s.Model != live.Model || s.ProviderName != live.ProviderName { + t.Errorf("absent config must inherit model/provider, got %q/%q", s.Model, s.ProviderName) + } + if s.ThinkingLevel != agentcore.ThinkingMedium { + t.Errorf("absent config must inherit thinkingLevel, got %q", s.ThinkingLevel) + } + if warn.Len() != 0 { + t.Errorf("absent config must not warn, got %q", warn.String()) + } +} + +// TestBtwConfigEmptyObjectInherits verifies that an empty JSON object inherits +// everything without warning. +func TestBtwConfigEmptyObjectInherits(t *testing.T) { + withBtwConfig(t, "{}") + live := &cli.LiveConfig{Model: "sess-model", ThinkingLevel: agentcore.ThinkingLow} + host := fakeHost{live: live} + + var warn bytes.Buffer + s := ResolveBtwSettings(&warn, host) + if s.Model != live.Model || s.ThinkingLevel != agentcore.ThinkingLow { + t.Errorf("empty object must inherit, got model=%q thinking=%q", s.Model, s.ThinkingLevel) + } + if warn.Len() != 0 { + t.Errorf("empty object must not warn, got %q", warn.String()) + } +} + +// TestBtwConfigThinkingOverride verifies a valid thinkingLevel is applied while +// the model still inherits (partial config falls back per-field). +func TestBtwConfigThinkingOverride(t *testing.T) { + withBtwConfig(t, `{"thinkingLevel":"high"}`) + live := &cli.LiveConfig{Model: "sess-model", ThinkingLevel: agentcore.ThinkingLow} + host := fakeHost{live: live} + + var warn bytes.Buffer + s := ResolveBtwSettings(&warn, host) + if s.ThinkingLevel != agentcore.ThinkingHigh { + t.Errorf("expected thinkingLevel override 'high', got %q", s.ThinkingLevel) + } + if s.Model != live.Model { + t.Errorf("model must still inherit when only thinkingLevel is set, got %q", s.Model) + } + if warn.Len() != 0 { + t.Errorf("valid override must not warn, got %q", warn.String()) + } +} + +// TestBtwConfigInvalidThinkingWarnsAndFallsBack verifies an invalid thinkingLevel +// warns on one line and keeps the session value. +func TestBtwConfigInvalidThinkingWarnsAndFallsBack(t *testing.T) { + withBtwConfig(t, `{"thinkingLevel":"bogus"}`) + live := &cli.LiveConfig{Model: "sess-model", ThinkingLevel: agentcore.ThinkingMedium} + host := fakeHost{live: live} + + var warn bytes.Buffer + s := ResolveBtwSettings(&warn, host) + if s.ThinkingLevel != agentcore.ThinkingMedium { + t.Errorf("invalid thinkingLevel must fall back to session value, got %q", s.ThinkingLevel) + } + if !strings.Contains(warn.String(), "thinkingLevel") { + t.Errorf("expected a warning about the invalid thinkingLevel, got %q", warn.String()) + } +} + +// TestBtwConfigMalformedWarnsAndInherits verifies a malformed JSON file warns +// once and inherits every field (never crashes /btw). +func TestBtwConfigMalformedWarnsAndInherits(t *testing.T) { + withBtwConfig(t, `{not json`) + live := &cli.LiveConfig{Model: "sess-model", ThinkingLevel: agentcore.ThinkingLow} + host := fakeHost{live: live} + + var warn bytes.Buffer + s := ResolveBtwSettings(&warn, host) + if s.Model != live.Model || s.ThinkingLevel != agentcore.ThinkingLow { + t.Errorf("malformed config must inherit, got model=%q thinking=%q", s.Model, s.ThinkingLevel) + } + if !strings.Contains(warn.String(), "invalid btw.json") { + t.Errorf("expected a malformed-config warning, got %q", warn.String()) + } +} + +// TestBtwConfigDoesNotMutateSession verifies ResolveBtwSettings never mutates +// the session live config, so a /btw override cannot leak into the main session +// (FR-8). +func TestBtwConfigDoesNotMutateSession(t *testing.T) { + withBtwConfig(t, `{"thinkingLevel":"xhigh"}`) + live := &cli.LiveConfig{Model: "sess-model", ThinkingLevel: agentcore.ThinkingLow} + host := fakeHost{live: live} + + _ = ResolveBtwSettings(&bytes.Buffer{}, host) + if live.ThinkingLevel != agentcore.ThinkingLow { + t.Errorf("session thinkingLevel must be unchanged by /btw config, got %q", live.ThinkingLevel) + } +} diff --git a/pigo/internal/cli/config/config.go b/pigo/internal/cli/config/config.go new file mode 100644 index 0000000..e83d1f4 --- /dev/null +++ b/pigo/internal/cli/config/config.go @@ -0,0 +1,128 @@ +// Package config implements pigo's optional user config file at +// ~/.config/pigo/config.toml (honoring $XDG_CONFIG_HOME when set) plus the +// provider-agnostic base-url env-var name derivation. Values in the file +// replace pigo's built-in defaults, but an explicit command-line flag always +// wins over the file: +// +// command-line flag > config.toml > built-in default +// +// A missing file is not an error (defaults apply); a malformed file is surfaced +// to the caller so it can warn rather than silently ignore user intent. +// +// The package is intentionally free of any cliOptions/run-assembly concern: it +// only loads and decodes the file and derives env-var names. Overlaying a +// FileConfig onto the parsed CLI options lives in cmd/pigo, alongside the +// options struct it mutates. +package config + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/BurntSushi/toml" +) + +// FileConfig is the on-disk shape of config.toml. Every field is optional; an +// absent (zero-value) field leaves the corresponding default/flag untouched. +// Keys are snake_case to read naturally in TOML. +type FileConfig struct { + Model string `toml:"model"` + BaseURL string `toml:"base_url"` + APIKey string `toml:"api_key"` + Protocol string `toml:"protocol"` + Provider string `toml:"provider"` + ThinkingLevel string `toml:"thinking_level"` + OutputFormat string `toml:"output_format"` + NoTools bool `toml:"no_tools"` + NoSkills bool `toml:"no_skills"` + Approve bool `toml:"approve"` + SystemPrompt string `toml:"system_prompt"` + // AllowedTools and DisallowedTools are the tool-level admission boundary: + // the config-file tier of --allowed-tools / --disallowed-tools. Names match + // case-insensitively and DisallowedTools wins when a name appears in both. + // A CLI flag replaces the file value wholesale rather than merging with it, + // so passing --allowed-tools can widen a boundary the file narrowed. + AllowedTools []string `toml:"allowed_tools"` + DisallowedTools []string `toml:"disallowed_tools"` + // Prompts is the config.toml `prompts` array: paths (files or dirs) to load + // prompt templates from at the settings tier (mirrors pi's settings prompts). + Prompts []string `toml:"prompts"` + // Memory, Checkpoint, and Compaction are nested TOML tables for the + // persistent-memory / infinite-context feature. They are pure config + // plumbing here; defaults/parsing live in memory.go (Resolve* helpers) and + // the overlay into runtime options lives in cmd/pigo. See + // tasks/spec-persistent-memory-infinite-context.md §3/§4/§5.2. + Memory MemoryConfig `toml:"memory"` + Checkpoint CheckpointConfig `toml:"checkpoint"` + Compaction CompactionConfig `toml:"compaction"` + // Dream is the [dream] TOML table for the /dream memory-consolidation + // feature. Pure config plumbing here; defaults/normalization live in + // internal/dream (Config). See tasks/spec-dream-memory-consolidation.md + // §3.3. + Dream DreamConfig `toml:"dream"` +} + +// DreamConfig is the [dream] TOML table for /dream memory consolidation. +// Enabled is a pointer so an absent key (nil) is distinguishable from an +// explicit false: nil is treated as true, only enabled = false disables +// auto-trigger. IntervalDays and RecentSessions use zero as "apply default" +// (7 and 20 respectively); normalization lives in dream.Config. +type DreamConfig struct { + Enabled *bool `toml:"enabled"` + IntervalDays int `toml:"interval_days"` + RecentSessions int `toml:"recent_sessions"` +} + +// FileConfigPath returns the path to the user config file: +// $XDG_CONFIG_HOME/pigo/config.toml, or ~/.config/pigo/config.toml by default. +// It returns "" when neither can be resolved, so the caller treats the file as +// absent. +func FileConfigPath() string { + if dir := os.Getenv("XDG_CONFIG_HOME"); dir != "" { + return filepath.Join(dir, "pigo", "config.toml") + } + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".config", "pigo", "config.toml") +} + +// LoadFileConfig reads and decodes config.toml. A missing file (or an empty +// path) returns a zero config with no error; a malformed file is an error. +func LoadFileConfig(path string) (FileConfig, error) { + if path == "" { + return FileConfig{}, nil + } + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return FileConfig{}, nil + } + return FileConfig{}, fmt.Errorf("read config %s: %w", path, err) + } + var cfg FileConfig + if err := toml.Unmarshal(data, &cfg); err != nil { + return FileConfig{}, fmt.Errorf("parse config %s: %w", path, err) + } + return cfg, nil +} + +// GenericBaseURLEnvVar derives the generic base-url override env var name for a +// provider: the provider name uppercased with hyphens rewritten to underscores, +// suffixed with _BASE_URL. For example "zai-coding-cn" → "ZAI_CODING_CN_BASE_URL" +// and "deepseek" → "DEEPSEEK_BASE_URL". An empty provider name yields "". +// +// It lives here (not with ResolveBaseURL) because it is a pure name derivation +// with no dependency on the provider registry — the provider-agnostic part of +// base-url resolution. ResolveBaseURL itself lives in internal/provider. +func GenericBaseURLEnvVar(providerName string) string { + n := strings.TrimSpace(providerName) + if n == "" { + return "" + } + n = strings.ReplaceAll(n, "-", "_") + return strings.ToUpper(n) + "_BASE_URL" +} diff --git a/pigo/internal/cli/config/config_test.go b/pigo/internal/cli/config/config_test.go new file mode 100644 index 0000000..517dcfb --- /dev/null +++ b/pigo/internal/cli/config/config_test.go @@ -0,0 +1,176 @@ +package config + +import ( + "os" + "path/filepath" + "reflect" + "testing" +) + +func TestFileConfigPath_XDGOverride(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", "/tmp/xdgroot") + got := FileConfigPath() + want := filepath.Join("/tmp/xdgroot", "pigo", "config.toml") + if got != want { + t.Fatalf("FileConfigPath() = %q, want %q", got, want) + } +} + +func TestFileConfigPath_DefaultHome(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", "") + home, err := os.UserHomeDir() + if err != nil { + t.Skip("no home dir") + } + got := FileConfigPath() + want := filepath.Join(home, ".config", "pigo", "config.toml") + if got != want { + t.Fatalf("FileConfigPath() = %q, want %q", got, want) + } +} + +func TestLoadFileConfig_Missing(t *testing.T) { + cfg, err := LoadFileConfig(filepath.Join(t.TempDir(), "does-not-exist.toml")) + if err != nil { + t.Fatalf("missing file should not error, got %v", err) + } + if !reflect.DeepEqual(cfg, FileConfig{}) { + t.Fatalf("missing file should yield zero config, got %+v", cfg) + } +} + +func TestLoadFileConfig_EmptyPath(t *testing.T) { + cfg, err := LoadFileConfig("") + if err != nil { + t.Fatalf("empty path should not error, got %v", err) + } + if !reflect.DeepEqual(cfg, FileConfig{}) { + t.Fatalf("empty path should yield zero config, got %+v", cfg) + } +} + +func TestLoadFileConfig_Valid(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + content := ` +model = "claude-opus-4-8" +base_url = "https://example.com" +api_key = "sk-test" +protocol = "anthropic" +provider = "deepseek" +thinking_level = "high" +output_format = "stream-json" +no_tools = true +no_skills = true +approve = true +system_prompt = "be terse" +` + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFileConfig(path) + if err != nil { + t.Fatalf("valid file should parse, got %v", err) + } + want := FileConfig{ + Model: "claude-opus-4-8", + BaseURL: "https://example.com", + APIKey: "sk-test", + Protocol: "anthropic", + Provider: "deepseek", + ThinkingLevel: "high", + OutputFormat: "stream-json", + NoTools: true, + NoSkills: true, + Approve: true, + SystemPrompt: "be terse", + } + if !reflect.DeepEqual(cfg, want) { + t.Fatalf("parsed config = %+v, want %+v", cfg, want) + } +} + +func TestLoadFileConfig_Malformed(t *testing.T) { + path := filepath.Join(t.TempDir(), "bad.toml") + if err := os.WriteFile(path, []byte("model = = ="), 0o600); err != nil { + t.Fatal(err) + } + if _, err := LoadFileConfig(path); err == nil { + t.Fatal("malformed file should error") + } +} + +func TestLoadFileConfigPromptsArray(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + content := "prompts = [\"./my-prompts\", \"/abs/x.md\"]\n" + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFileConfig(path) + if err != nil { + t.Fatalf("LoadFileConfig: %v", err) + } + if len(cfg.Prompts) != 2 || cfg.Prompts[0] != "./my-prompts" || cfg.Prompts[1] != "/abs/x.md" { + t.Errorf("Prompts = %v, want [./my-prompts /abs/x.md]", cfg.Prompts) + } +} + +// TestGenericBaseURLEnvVar verifies the _BASE_URL name derivation, +// especially the hyphen→underscore conversion and uppercasing. +func TestGenericBaseURLEnvVar(t *testing.T) { + cases := []struct { + name string + want string + }{ + {"deepseek", "DEEPSEEK_BASE_URL"}, + {"zai-coding-cn", "ZAI_CODING_CN_BASE_URL"}, + {"vercel-ai-gateway", "VERCEL_AI_GATEWAY_BASE_URL"}, + {"", ""}, + } + for _, c := range cases { + if got := GenericBaseURLEnvVar(c.name); got != c.want { + t.Errorf("GenericBaseURLEnvVar(%q) = %q, want %q", c.name, got, c.want) + } + } +} + +func TestLoadFileConfig_DreamTable(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + content := ` +[dream] +enabled = false +interval_days = 14 +recent_sessions = 50 +` + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFileConfig(path) + if err != nil { + t.Fatalf("LoadFileConfig: %v", err) + } + if cfg.Dream.Enabled == nil || *cfg.Dream.Enabled { + t.Errorf("Dream.Enabled = %v, want explicit false", cfg.Dream.Enabled) + } + if cfg.Dream.IntervalDays != 14 { + t.Errorf("Dream.IntervalDays = %d, want 14", cfg.Dream.IntervalDays) + } + if cfg.Dream.RecentSessions != 50 { + t.Errorf("Dream.RecentSessions = %d, want 50", cfg.Dream.RecentSessions) + } +} + +func TestLoadFileConfig_DreamTableAbsent(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + if err := os.WriteFile(path, []byte("model = \"foo\"\n"), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFileConfig(path) + if err != nil { + t.Fatalf("LoadFileConfig: %v", err) + } + // Absent [dream] table: Enabled pointer nil (→ default true downstream), + // ints zero (→ defaults downstream). Parsing must not error. + if cfg.Dream.Enabled != nil || cfg.Dream.IntervalDays != 0 || cfg.Dream.RecentSessions != 0 { + t.Errorf("absent dream table = %+v, want zero-value", cfg.Dream) + } +} diff --git a/pigo/internal/cli/config/memory.go b/pigo/internal/cli/config/memory.go new file mode 100644 index 0000000..0e2da77 --- /dev/null +++ b/pigo/internal/cli/config/memory.go @@ -0,0 +1,265 @@ +// Memory/checkpoint/compaction config: the nested TOML tables [memory], +// [checkpoint], and compaction.max_context (spec-persistent-memory-infinite- +// context §3/§4/§5.2). This file is pure config plumbing — parse, defaults, and +// resolve helpers only. The actual memory Store, checkpoint persistence, and +// compaction trigger live in later layers (internal/memory, internal/runtime, +// internal/compaction) and consume the resolved values overlaid in cmd/pigo. +package config + +import ( + "fmt" + "strconv" + "strings" +) + +// Built-in defaults for the [memory] table. Exposed so overlay and tests share +// one source of truth (mirrors the flat-config defaults documented in main.go). +const ( + DefaultMemoryEnabled = true + DefaultMemoryReconcileOnSearch = true + DefaultMemorySearchScoreFloor = 0.15 + DefaultMemoryCCIndex = false +) + +// DefaultCheckpointThresholds is the built-in compaction trigger ladder as +// percentage strings; ResolveThresholds parses them into fractions in (0,1]. +func DefaultCheckpointThresholds() []string { + return []string{"40%", "60%", "80%"} +} + +// MemoryConfig is the [memory] TOML table. Bool fields whose default is true +// (Enabled, ReconcileOnSearch) and the float ScoreFloor use pointers so an +// absent key is distinguishable from an explicit false/0 — nil means "apply the +// default", which is what makes memory.enabled=false representable and +// default-safe. CCIndex defaults to false, so a plain bool suffices. +type MemoryConfig struct { + Enabled *bool `toml:"enabled"` + ReconcileOnSearch *bool `toml:"reconcile_on_search"` + SearchScoreFloor *float64 `toml:"search_score_floor"` + CCIndex bool `toml:"cc_index"` +} + +// ResolvedMemory is MemoryConfig with defaults applied and the score floor +// clamped to [0,1]. It is the shape downstream memory wiring consumes. +type ResolvedMemory struct { + Enabled bool + ReconcileOnSearch bool + SearchScoreFloor float64 + CCIndex bool +} + +// Resolve applies the [memory] defaults: absent keys fall back to +// true/true/0.15/false; an explicit search_score_floor outside [0,1] is clamped +// into range. +func (m MemoryConfig) Resolve() ResolvedMemory { + r := ResolvedMemory{ + Enabled: DefaultMemoryEnabled, + ReconcileOnSearch: DefaultMemoryReconcileOnSearch, + SearchScoreFloor: DefaultMemorySearchScoreFloor, + CCIndex: m.CCIndex, + } + if m.Enabled != nil { + r.Enabled = *m.Enabled + } + if m.ReconcileOnSearch != nil { + r.ReconcileOnSearch = *m.ReconcileOnSearch + } + if m.SearchScoreFloor != nil { + f := *m.SearchScoreFloor + switch { + case f < 0: + f = 0 + case f > 1: + f = 1 + } + r.SearchScoreFloor = f + } + return r +} + +// IntOrString accepts either a TOML integer or string for the optional +// [checkpoint].reserved key (e.g. reserved = 4096 or reserved = "10%"). Set +// reports whether the key was present; IsInt selects the populated field. +type IntOrString struct { + Set bool + IsInt bool + Int int + Str string +} + +// UnmarshalTOML implements toml.Unmarshaler so a bare int or a quoted string +// both decode without failing the whole file. +func (v *IntOrString) UnmarshalTOML(data any) error { + v.Set = true + switch t := data.(type) { + case int64: + v.IsInt, v.Int = true, int(t) + case int: + v.IsInt, v.Int = true, t + case float64: + v.IsInt, v.Int = true, int(t) + case string: + v.Str = t + default: + return fmt.Errorf("reserved: unsupported type %T (want int or string)", data) + } + return nil +} + +// CheckpointConfig is the [checkpoint] TOML table. push_caps is a nested table +// of per-section token caps (e.g. [checkpoint.push_caps] with memory = 800, +// recall = 1200), modeled as a map. +type CheckpointConfig struct { + Thresholds []string `toml:"thresholds"` + Reserved IntOrString `toml:"reserved"` + PushCaps map[string]int `toml:"push_caps"` +} + +// ResolveThresholds parses the configured threshold percentage strings into +// fractions in (0,1], skipping out-of-range/unparseable entries. An empty list +// — or one where every entry is invalid — falls back to the built-in defaults. +func (c CheckpointConfig) ResolveThresholds() []float64 { + src := c.Thresholds + if len(src) == 0 { + src = DefaultCheckpointThresholds() + } + out := parseThresholds(src) + if len(out) == 0 { + out = parseThresholds(DefaultCheckpointThresholds()) + } + return out +} + +func parseThresholds(ss []string) []float64 { + var out []float64 + for _, s := range ss { + if f, ok := ParseThresholdFraction(s); ok { + out = append(out, f) + } + } + return out +} + +// ParseThresholdFraction parses a percentage string like "80%" into a fraction +// in (0,1]. Values outside that range (<=0, >100%) or lacking a % suffix are +// rejected with ok=false so callers fall back to a default. +func ParseThresholdFraction(s string) (float64, bool) { + s = strings.TrimSpace(s) + if !strings.HasSuffix(s, "%") { + return 0, false + } + n, err := strconv.ParseFloat(strings.TrimSpace(strings.TrimSuffix(s, "%")), 64) + if err != nil { + return 0, false + } + f := n / 100 + if f <= 0 || f > 1 { + return 0, false + } + return f, true +} + +// CompactionConfig is the [compaction] TOML table. Only max_context is wired +// here; it lowers the auto-compaction trigger point and is always clamped by +// the provider window by the consumer. +type CompactionConfig struct { + MaxContext string `toml:"max_context"` +} + +// ResolveMaxContext parses the max_context string form. An empty value yields +// an unset MaxContext (no error). +func (c CompactionConfig) ResolveMaxContext() (MaxContext, error) { + return ParseMaxContext(c.MaxContext) +} + +// MaxContext is a parsed compaction.max_context value: either an absolute token +// count or a fraction of the provider window. The zero value is "unset" and +// Resolve returns 0. Resolve is a pure function; the provider-limit clamp is +// applied by the consumer. +type MaxContext struct { + set bool + fraction float64 // >0 for a "N%" form + tokens int // absolute token count when fraction == 0 +} + +// IsSet reports whether max_context was configured. +func (m MaxContext) IsSet() bool { return m.set } + +// Resolve returns the token budget for the given provider window: window* +// fraction (rounded) for a percentage form, or the absolute token count +// otherwise. An unset value returns 0. This is intentionally unclamped — the +// consumer applies the provider-limit clamp. +func (m MaxContext) Resolve(window int) int { + if !m.set { + return 0 + } + if m.fraction > 0 { + return int(float64(window)*m.fraction + 0.5) + } + return m.tokens +} + +// ParseMaxContext parses the accepted max_context forms: a plain token count +// ("300000"), a K/M-suffixed count ("300K", "1M", case-insensitive, fractions +// allowed like "1.5M"), or a percentage of the provider window ("50%"). An +// empty string is unset (no error); other malformed or non-positive values are +// errors. +func ParseMaxContext(s string) (MaxContext, error) { + s = strings.TrimSpace(s) + if s == "" { + return MaxContext{}, nil + } + if strings.HasSuffix(s, "%") { + n, err := strconv.ParseFloat(strings.TrimSpace(strings.TrimSuffix(s, "%")), 64) + if err != nil { + return MaxContext{}, fmt.Errorf("max_context: invalid percent %q: %w", s, err) + } + f := n / 100 + if f <= 0 || f > 1 { + return MaxContext{}, fmt.Errorf("max_context: percent out of range %q (want (0%%,100%%])", s) + } + return MaxContext{set: true, fraction: f}, nil + } + mult := 1.0 + body := s + switch last := s[len(s)-1]; last { + case 'k', 'K': + mult, body = 1_000, s[:len(s)-1] + case 'm', 'M': + mult, body = 1_000_000, s[:len(s)-1] + } + n, err := strconv.ParseFloat(strings.TrimSpace(body), 64) + if err != nil { + return MaxContext{}, fmt.Errorf("max_context: invalid token count %q: %w", s, err) + } + if n <= 0 { + return MaxContext{}, fmt.Errorf("max_context: must be positive %q", s) + } + return MaxContext{set: true, tokens: int(n*mult + 0.5)}, nil +} + +// MemorySettings bundles the resolved [memory]/[checkpoint]/[compaction] config +// for overlay into runtime options: defaults applied, string forms pre-parsed. +// It is produced by FileConfig.ResolveMemorySettings and is always well-formed +// (an invalid max_context is treated as unset rather than failing the overlay). +type MemorySettings struct { + Memory ResolvedMemory + CheckpointThresholds []float64 + CheckpointReserved IntOrString + CheckpointPushCaps map[string]int + MaxContext MaxContext +} + +// ResolveMemorySettings resolves the three nested tables into MemorySettings, +// applying defaults. It never fails: an unparseable compaction.max_context is +// dropped to unset so a single bad key cannot break config overlay. +func (c FileConfig) ResolveMemorySettings() MemorySettings { + mc, _ := c.Compaction.ResolveMaxContext() + return MemorySettings{ + Memory: c.Memory.Resolve(), + CheckpointThresholds: c.Checkpoint.ResolveThresholds(), + CheckpointReserved: c.Checkpoint.Reserved, + CheckpointPushCaps: c.Checkpoint.PushCaps, + MaxContext: mc, + } +} diff --git a/pigo/internal/cli/config/memory_test.go b/pigo/internal/cli/config/memory_test.go new file mode 100644 index 0000000..f2438e1 --- /dev/null +++ b/pigo/internal/cli/config/memory_test.go @@ -0,0 +1,234 @@ +package config + +import ( + "os" + "path/filepath" + "reflect" + "testing" +) + +// boolPtr / floatPtr are test helpers for the pointer-valued MemoryConfig fields. +func boolPtr(b bool) *bool { return &b } +func floatPtr(f float64) *float64 { return &f } + +func TestMemoryConfig_ResolveDefaults(t *testing.T) { + got := MemoryConfig{}.Resolve() + want := ResolvedMemory{ + Enabled: true, + ReconcileOnSearch: true, + SearchScoreFloor: 0.15, + CCIndex: false, + } + if got != want { + t.Fatalf("Resolve() defaults = %+v, want %+v", got, want) + } +} + +func TestMemoryConfig_ResolveOverrides(t *testing.T) { + got := MemoryConfig{ + Enabled: boolPtr(false), + ReconcileOnSearch: boolPtr(false), + SearchScoreFloor: floatPtr(0.5), + CCIndex: true, + }.Resolve() + want := ResolvedMemory{ + Enabled: false, + ReconcileOnSearch: false, + SearchScoreFloor: 0.5, + CCIndex: true, + } + if got != want { + t.Fatalf("Resolve() overrides = %+v, want %+v", got, want) + } +} + +// memory.enabled=false must be representable and distinct from "absent". +func TestMemoryConfig_EnabledFalseRepresentable(t *testing.T) { + if r := (MemoryConfig{Enabled: boolPtr(false)}).Resolve(); r.Enabled { + t.Fatal("explicit enabled=false should resolve to false") + } + if r := (MemoryConfig{}).Resolve(); !r.Enabled { + t.Fatal("absent enabled should default to true") + } +} + +func TestMemoryConfig_ScoreFloorClamp(t *testing.T) { + if r := (MemoryConfig{SearchScoreFloor: floatPtr(-1)}).Resolve(); r.SearchScoreFloor != 0 { + t.Errorf("negative score floor should clamp to 0, got %v", r.SearchScoreFloor) + } + if r := (MemoryConfig{SearchScoreFloor: floatPtr(2)}).Resolve(); r.SearchScoreFloor != 1 { + t.Errorf("score floor >1 should clamp to 1, got %v", r.SearchScoreFloor) + } +} + +func TestParseThresholdFraction(t *testing.T) { + cases := []struct { + in string + want float64 + ok bool + }{ + {"40%", 0.40, true}, + {"60%", 0.60, true}, + {"80%", 0.80, true}, + {"100%", 1.0, true}, + {" 75% ", 0.75, true}, + {"0%", 0, false}, // out of range (must be >0) + {"120%", 0, false}, // out of range (>100%) + {"-10%", 0, false}, + {"80", 0, false}, // missing % + {"abc%", 0, false}, + } + for _, c := range cases { + got, ok := ParseThresholdFraction(c.in) + if ok != c.ok || (ok && got != c.want) { + t.Errorf("ParseThresholdFraction(%q) = (%v,%v), want (%v,%v)", c.in, got, ok, c.want, c.ok) + } + } +} + +func TestCheckpointConfig_ResolveThresholds(t *testing.T) { + // Absent → defaults 40/60/80%. + if got := (CheckpointConfig{}).ResolveThresholds(); !reflect.DeepEqual(got, []float64{0.40, 0.60, 0.80}) { + t.Errorf("default thresholds = %v, want [0.4 0.6 0.8]", got) + } + // Explicit override. + if got := (CheckpointConfig{Thresholds: []string{"50%", "90%"}}).ResolveThresholds(); !reflect.DeepEqual(got, []float64{0.50, 0.90}) { + t.Errorf("override thresholds = %v, want [0.5 0.9]", got) + } + // Out-of-range entries skipped, valid kept. + if got := (CheckpointConfig{Thresholds: []string{"120%", "70%"}}).ResolveThresholds(); !reflect.DeepEqual(got, []float64{0.70}) { + t.Errorf("mixed thresholds = %v, want [0.7]", got) + } + // All invalid → fall back to defaults. + if got := (CheckpointConfig{Thresholds: []string{"nope", "0%"}}).ResolveThresholds(); !reflect.DeepEqual(got, []float64{0.40, 0.60, 0.80}) { + t.Errorf("all-invalid thresholds = %v, want defaults", got) + } +} + +func TestParseMaxContext(t *testing.T) { + cases := []struct { + in string + set bool + window int + want int + }{ + {"", false, 200000, 0}, + {"300000", true, 200000, 300000}, + {"300K", true, 200000, 300000}, + {"1M", true, 200000, 1000000}, + {"1m", true, 200000, 1000000}, + {"1.5M", true, 200000, 1500000}, + {"50%", true, 200000, 100000}, + {"25%", true, 400000, 100000}, + } + for _, c := range cases { + mc, err := ParseMaxContext(c.in) + if err != nil { + t.Errorf("ParseMaxContext(%q) unexpected error: %v", c.in, err) + continue + } + if mc.IsSet() != c.set { + t.Errorf("ParseMaxContext(%q).IsSet() = %v, want %v", c.in, mc.IsSet(), c.set) + } + if got := mc.Resolve(c.window); got != c.want { + t.Errorf("ParseMaxContext(%q).Resolve(%d) = %d, want %d", c.in, c.window, got, c.want) + } + } +} + +func TestParseMaxContext_Invalid(t *testing.T) { + for _, in := range []string{"0%", "150%", "-100%", "abc", "0", "-5", "12x"} { + if _, err := ParseMaxContext(in); err == nil { + t.Errorf("ParseMaxContext(%q) should error", in) + } + } +} + +func TestLoadFileConfig_MemoryTables(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + content := ` +[memory] +enabled = false +reconcile_on_search = false +search_score_floor = 0.3 +cc_index = true + +[checkpoint] +thresholds = ["50%", "70%"] +reserved = 4096 + +[checkpoint.push_caps] +memory = 800 +recall = 1200 + +[compaction] +max_context = "300K" +` + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFileConfig(path) + if err != nil { + t.Fatalf("LoadFileConfig: %v", err) + } + + mem := cfg.Memory.Resolve() + if mem.Enabled || mem.ReconcileOnSearch || !mem.CCIndex || mem.SearchScoreFloor != 0.3 { + t.Errorf("memory resolved = %+v", mem) + } + if got := cfg.Checkpoint.ResolveThresholds(); !reflect.DeepEqual(got, []float64{0.50, 0.70}) { + t.Errorf("thresholds = %v, want [0.5 0.7]", got) + } + if !cfg.Checkpoint.Reserved.Set || !cfg.Checkpoint.Reserved.IsInt || cfg.Checkpoint.Reserved.Int != 4096 { + t.Errorf("reserved = %+v, want int 4096", cfg.Checkpoint.Reserved) + } + if cfg.Checkpoint.PushCaps["memory"] != 800 || cfg.Checkpoint.PushCaps["recall"] != 1200 { + t.Errorf("push_caps = %v", cfg.Checkpoint.PushCaps) + } + mc, err := cfg.Compaction.ResolveMaxContext() + if err != nil { + t.Fatalf("ResolveMaxContext: %v", err) + } + if got := mc.Resolve(200000); got != 300000 { + t.Errorf("max_context resolve = %d, want 300000", got) + } +} + +func TestLoadFileConfig_ReservedString(t *testing.T) { + path := filepath.Join(t.TempDir(), "config.toml") + if err := os.WriteFile(path, []byte("[checkpoint]\nreserved = \"10%\"\n"), 0o600); err != nil { + t.Fatal(err) + } + cfg, err := LoadFileConfig(path) + if err != nil { + t.Fatalf("LoadFileConfig: %v", err) + } + if !cfg.Checkpoint.Reserved.Set || cfg.Checkpoint.Reserved.IsInt || cfg.Checkpoint.Reserved.Str != "10%" { + t.Errorf("reserved = %+v, want string 10%%", cfg.Checkpoint.Reserved) + } +} + +// Absent tables → default-safe resolved settings. +func TestResolveMemorySettings_Defaults(t *testing.T) { + ms := FileConfig{}.ResolveMemorySettings() + if !ms.Memory.Enabled || !ms.Memory.ReconcileOnSearch || ms.Memory.SearchScoreFloor != 0.15 || ms.Memory.CCIndex { + t.Errorf("default memory = %+v", ms.Memory) + } + if !reflect.DeepEqual(ms.CheckpointThresholds, []float64{0.40, 0.60, 0.80}) { + t.Errorf("default thresholds = %v", ms.CheckpointThresholds) + } + if ms.MaxContext.IsSet() { + t.Errorf("default max_context should be unset") + } + if ms.CheckpointReserved.Set { + t.Errorf("default reserved should be unset") + } +} + +// An invalid max_context must not fail the overlay — it drops to unset. +func TestResolveMemorySettings_InvalidMaxContextIgnored(t *testing.T) { + ms := FileConfig{Compaction: CompactionConfig{MaxContext: "garbage"}}.ResolveMemorySettings() + if ms.MaxContext.IsSet() { + t.Errorf("invalid max_context should resolve to unset, got set") + } +} diff --git a/pigo/internal/cli/doc.go b/pigo/internal/cli/doc.go new file mode 100644 index 0000000..9584d97 --- /dev/null +++ b/pigo/internal/cli/doc.go @@ -0,0 +1,23 @@ +// Package cli is the contract and shared-type layer for pigo's command-line +// surface. It holds the interfaces and value types that the cmd/pigo entry +// point and the internal/cli/* subpackages share, without depending on any of +// them, so those subpackages can be assembled and tested in isolation. +// +// Subpackage layout (assembled incrementally by the CLI restructure): +// +// internal/cli — Host/Editor contracts, LiveConfig, TelemetryHolder +// internal/cli/ui — color, markdown, imageref helpers +// internal/cli/config — config-file loading and overrides +// internal/cli/run — run assembly and tool wiring +// internal/cli/headless — headless session and subagent RPC drivers +// internal/cli/goal — /goal state machine +// internal/cli/btw — /btw side thread +// internal/cli/status — /status command +// internal/cli/repl — interactive REPL and line editor +// internal/cli/pkgcmd — package-manager subcommands +// internal/cli/testutil — cross-subpackage test helpers +// +// The Host interface (see host.go) is the seam that lets the /goal, /btw, +// /status and REPL subpackages read and mutate the session's live state +// without importing the concrete replDeps aggregate that assembles it. +package cli diff --git a/pigo/internal/cli/goal/goal.go b/pigo/internal/cli/goal/goal.go new file mode 100644 index 0000000..3840b30 --- /dev/null +++ b/pigo/internal/cli/goal/goal.go @@ -0,0 +1,425 @@ +// This file implements the /goal command (mirrors pi-goal / Claude Code's goal +// mode): given a high-level objective, pigo runs the agent autonomously — +// re-prompting it turn after turn from the loop's follow-up seam — until the +// model declares the goal done (goal_complete), reports a true impasse +// (goal_blocked), or a safety guard (max turns / no-progress) or the token +// budget stops it. +// +// /goal is intercepted in the REPL loop rather than routed through a slash +// Action closure because it must run agent streams and mutate the shared +// context and goal state — none of which a pure string→string Action can do, +// exactly like /compact and /fork. It reaches the session's collaborators and +// mutable state through the cli.Host contract, so it need not import the +// concrete replDeps aggregate that assembles them. +// +// Scope: core autonomous continuation plus an optional token budget. Goal state +// lives only in the session's in-memory GoalState (host.Goal()); it is not +// persisted across process restarts. +package goal + +import ( + "context" + "fmt" + "io" + "strconv" + "strings" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/trust" +) + +// goalMaxAutomaticTurns caps how many autonomous continuations a single /goal +// run will issue before pausing, a runaway guard (mirrors pi-goal's automaticTurns). +const goalMaxAutomaticTurns = 25 + +// goalMaxNoProgress pauses the run after this many consecutive tool-free +// continuations, so a model looping without acting cannot spin forever (mirrors +// pi-goal's noProgressTurns). +const goalMaxNoProgress = 3 + +// runGoal parses and dispatches a /goal invocation. setCancel publishes the +// active run's cancel func so the REPL's SIGINT handler can interrupt an +// autonomous run (same plumbing as a normal turn). +func RunGoal(setCancel func(context.CancelFunc), out io.Writer, host cli.Host, line string) { + args := strings.TrimSpace(strings.TrimPrefix(line, "/goal")) + + switch { + case args == "": + printGoalStatus(out, host.Goal()) + return + case args == "clear": + host.Goal().Clear() + fmt.Fprintln(out, "goal cleared") + return + case args == "pause": + snap := host.Goal().Snapshot() + if snap.Status != agenttool.GoalActive && snap.Status != agenttool.GoalPaused { + fmt.Fprintln(out, "no active goal to pause") + return + } + host.Goal().SetStatus(agenttool.GoalPaused) + fmt.Fprintln(out, "goal paused — run /goal resume to continue") + return + case args == "resume": + snap := host.Goal().Snapshot() + if snap.Status != agenttool.GoalPaused && snap.Status != agenttool.GoalBudgetLimited { + fmt.Fprintln(out, "no paused goal to resume") + return + } + host.Goal().Resume() + fmt.Fprintf(out, "resuming goal: %s\n", ui.OneLine(snap.Objective)) + runGoalLoop(setCancel, out, host) + return + } + + // Otherwise args is a new objective, optionally prefixed with --tokens N. + objective, budget, err := parseGoalObjective(args) + if err != nil { + fmt.Fprintf(out, "pigo: %v\n", err) + return + } + if strings.TrimSpace(objective) == "" { + fmt.Fprintln(out, "usage: /goal [--tokens N] ") + return + } + host.Goal().Start(newGoalID(), objective, budget) + if budget > 0 { + fmt.Fprintf(out, "goal set (token budget %d): %s\n", budget, ui.OneLine(objective)) + } else { + fmt.Fprintf(out, "goal set: %s\n", ui.OneLine(objective)) + } + runGoalLoop(setCancel, out, host) +} + +// newGoalID returns a short unique id for a goal (used to key goal_complete's +// exact-id contract in a future extension; here it just labels the goal). +func newGoalID() string { return "goal-" + strconv.FormatInt(time.Now().UnixNano(), 36) } + +// parseGoalObjective splits an optional leading "--tokens N" flag from the +// objective text. N accepts a bare integer or a k/m suffix (100k, 1m). The flag +// must lead; anything after it (or the whole string when absent) is the +// objective. +func parseGoalObjective(args string) (objective string, budget int, err error) { + rest := args + if strings.HasPrefix(rest, "--tokens") { + rest = strings.TrimSpace(strings.TrimPrefix(rest, "--tokens")) + // The value is the next whitespace-delimited token. + var valTok string + if i := strings.IndexAny(rest, " \t"); i >= 0 { + valTok = rest[:i] + rest = strings.TrimSpace(rest[i+1:]) + } else { + valTok = rest + rest = "" + } + budget, err = parseTokenBudget(valTok) + if err != nil { + return "", 0, err + } + } + return rest, budget, nil +} + +// parseTokenBudget parses a token count with an optional k/m (×1000/×1000000) +// suffix, case-insensitive. It rejects a non-positive or malformed value. +func parseTokenBudget(s string) (int, error) { + s = strings.TrimSpace(strings.ToLower(s)) + if s == "" { + return 0, fmt.Errorf("--tokens requires a value (e.g. --tokens 100k)") + } + mult := 1 + switch { + case strings.HasSuffix(s, "k"): + mult = 1000 + s = strings.TrimSuffix(s, "k") + case strings.HasSuffix(s, "m"): + mult = 1000000 + s = strings.TrimSuffix(s, "m") + } + n, convErr := strconv.Atoi(s) + if convErr != nil || n <= 0 { + return 0, fmt.Errorf("invalid --tokens value %q (want a positive number, optionally with k/m)", s) + } + return n * mult, nil +} + +// goalContinuationPrompt is the follow-up injected each autonomous turn to keep +// the model working toward the goal. The objective itself is re-stated every +// turn by the GoalReminderProvider, so this only nudges continuation. +const goalContinuationPrompt = "Continue working toward the goal. When every requirement is " + + "verifiably met, call goal_complete with a summary. If you hit a true impasse you cannot work " + + "around, call goal_blocked with concrete evidence. Otherwise keep going — do not stop or ask " + + "the user whether to continue." + +// goalFollowUpDecision decides, from a goal snapshot, whether the autonomous +// loop should issue another continuation turn. It is pure so the branch logic +// (complete/blocked terminal states, token budget, turn cap, no-progress guard) +// can be unit-tested without running a real agent stream. The returned status is +// the terminal status to record when cont is false (GoalIdle means "leave the +// status as-is" — used for the already-terminal complete/blocked cases). +func goalFollowUpDecision(snap agenttool.GoalSnapshot) (cont bool, terminal agenttool.GoalStatus) { + switch snap.Status { + case agenttool.GoalComplete, agenttool.GoalBlocked: + // A goal tool already ended the run; nothing to continue. + return false, agenttool.GoalIdle + } + if snap.TokenBudget > 0 && snap.TokensUsed >= snap.TokenBudget { + return false, agenttool.GoalBudgetLimited + } + if snap.Iterations >= goalMaxAutomaticTurns { + return false, agenttool.GoalPaused + } + if snap.NoProgress >= goalMaxNoProgress { + return false, agenttool.GoalPaused + } + return true, agenttool.GoalIdle +} + +// runGoalLoop drives the autonomous goal run. It assembles a run just like +// streamRun but with the goal tools (goal_complete/goal_blocked) added, the goal +// reminder wired alongside the todo reminder, and a GetFollowUpMessages hook that +// re-prompts the model each time the inner loop settles — until a goal tool ends +// the run or goalFollowUpDecision trips a guard. On return it prints the outcome +// and persists the turn. The run reuses the REPL's SIGINT cancel plumbing via +// setCancel. +func runGoalLoop(setCancel func(context.CancelFunc), out io.Writer, host cli.Host) { + goalReg := goalToolRegistry(host.Registry(), host.Goal()) + reminders := goalReminders(host.Registry(), host.Goal()) + + runCtx, cancel := context.WithCancel(context.Background()) + setCancel(cancel) + defer func() { + cancel() + setCancel(nil) + }() + + // lastSeen tracks how many messages we have already accounted for, so each + // settle folds in only the assistant turns produced since the previous one: + // their output tokens (budget) and whether any tool ran (no-progress guard). + lastSeen := len(host.AgentCtx().Messages) + + cfg := runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: host.Live().Model, + Provider: host.Live().ProviderName, + ThinkingLevel: host.Live().ThinkingLevel, + Stream: provider.StreamFnFromProvider(host.Live().Provider), + GetAPIKey: host.Creds().GetAPIKey, + ContextWindow: host.Live().ContextWindow, + Compaction: compaction.DefaultCompactionSettings, + }, + Batch: agenttool.BatchConfig{ + ToolExecutorConfig: agenttool.ToolExecutorConfig{ + Registry: goalReg, + BeforeToolCall: trust.BeforeToolCall(host.Trust(), host.Cwd(), host.Input(), out, host.ConfirmMu()), + }, + }, + Reminders: reminders, + GetFollowUpMessages: func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage { + // Account for the turns produced since the last settle. Auto-compaction + // can shrink agentCtx.Messages in place (summary + tail) between settles, + // dropping its length below lastSeen; clamp so the slice never goes out + // of bounds. The compacted turn's tokens are then under-counted, which is + // acceptable for a soft budget guard (a crash is not). + if lastSeen > len(agentCtx.Messages) { + lastSeen = len(agentCtx.Messages) + } + outputTokens, hadTool := goalTurnActivity(agentCtx.Messages[lastSeen:]) + lastSeen = len(agentCtx.Messages) + host.Goal().RecordIteration(outputTokens, hadTool) + + cont, terminal := goalFollowUpDecision(host.Goal().Snapshot()) + if !cont { + if terminal != agenttool.GoalIdle { + host.Goal().SetStatus(terminal) + } + return nil + } + return []agentcore.AgentMessage{agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(goalContinuationPrompt)}, + }} + }, + } + + // Wire the per-turn hook seams (PreToolUse/PostToolUse/Stop) onto this goal + // run's cfg from the session dispatcher; a nil dispatcher is a no-op (FR-18). + if d := host.Dispatcher(); d != nil { + run.InstallSeams(&cfg, d, host.HookDeps()) + } + + // The first turn is driven by the goal reminder alone (the objective is + // injected as background context); no explicit user prompt is appended so the + // objective is not duplicated in the durable history. + stream := runtime.StartRun(runCtx, host.AgentCtx(), cfg) + drainGoalStream(runCtx, out, host, stream) + + printGoalOutcome(out, host.Goal().Snapshot()) + cli.PersistTurn(out, host) +} + +// goalToolRegistry returns a registry holding every tool from base plus the two +// goal-control tools, so the autonomous run can invoke goal_complete/goal_blocked +// while keeping all the normal tools available. The base registry is left +// unchanged (the goal tools are only present for the goal run). +func goalToolRegistry(base *agenttool.ToolRegistry, state *agenttool.GoalState) *agenttool.ToolRegistry { + reg := agenttool.NewToolRegistry() + if base != nil { + for _, t := range base.List() { + _ = reg.Register(t) + } + } + _ = reg.Register(&agenttool.GoalCompleteTool{State: state}) + _ = reg.Register(&agenttool.GoalBlockedTool{State: state}) + return reg +} + +// goalReminders builds the per-turn reminder registry for a goal run: the goal +// reminder (re-stating the objective every turn) plus the todo reminder when a +// todo tool is present, so an autonomous run keeps both its objective and its +// task list in view. +func goalReminders(base *agenttool.ToolRegistry, state *agenttool.GoalState) *runtime.ReminderRegistry { + reg := runtime.NewReminderRegistry(&runtime.GoalReminderProvider{State: state}) + if base != nil { + if t, ok := base.Get("todo"); ok { + if tt, ok := t.(*agenttool.TodoTool); ok && tt.Store != nil { + reg.Register(&runtime.TodoReminderProvider{Store: tt.Store}) + } + } + } + return reg +} + +// goalTurnActivity sums the output tokens across the assistant messages in tail +// and reports whether any tool ran in that window (an assistant tool call or a +// tool result). It feeds RecordIteration's token-budget and no-progress inputs. +func goalTurnActivity(tail []agentcore.AgentMessage) (outputTokens int, hadTool bool) { + for _, m := range tail { + switch msg := m.(type) { + case agentcore.AssistantMessage: + if msg.Usage != nil { + outputTokens += msg.Usage.OutputTokens + } + if len(msg.ToolCalls()) > 0 { + hadTool = true + } + case agentcore.ToolResultMessage: + hadTool = true + } + } + return outputTokens, hadTool +} + +// drainGoalStream prints the streamed assistant text and tool activity of a goal +// run to out, mirroring streamRun's rendering. It blocks until the run ends. +// chainGoalEvent returns the OnEvent observer for a goal run: the plugin +// notifier, with the SessionEnd/PreCompact hook notifier chained after it when +// hooks are configured, mirroring the REPL's OnEvent composition. +func chainGoalEvent(host cli.Host) func(agentcore.AgentEvent) { + notifier := host.NotifierHandle() + d := host.Dispatcher() + if d == nil { + return notifier + } + deps := host.HookDeps() + hookEvent := hooks.NewHookNotifier(d, deps.SessionID, deps.ProjectDir).Handle + if notifier == nil { + return hookEvent + } + return func(ev agentcore.AgentEvent) { + notifier(ev) + hookEvent(ev) + } +} + +func drainGoalStream(ctx context.Context, out io.Writer, host cli.Host, stream *runtime.LoopEventStream) { + var reply strings.Builder + flushReply := func() { + if reply.Len() == 0 { + return + } + rendered := ui.RenderMarkdown(reply.String()) + fmt.Fprint(out, rendered) + if !strings.HasSuffix(rendered, "\n") { + fmt.Fprintln(out) + } + reply.Reset() + } + _, err := runtime.DrainStream(ctx, stream, runtime.StreamHandler{ + OnEvent: chainGoalEvent(host), + OnText: func(delta string) { + reply.WriteString(delta) + }, + OnTurnEnd: func(msg agentcore.AssistantMessage, results []agentcore.ToolResultMessage) { + flushReply() + for _, c := range msg.ToolCalls() { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(ui.Enabled(), ui.Green, "→ tool:"), ui.ToolCallLabel(c)) + } + for _, tr := range results { + ui.RenderToolResult(out, tr) + } + }, + }) + flushReply() + if err != nil { + if ctx.Err() != nil { + fmt.Fprintln(out, "^C interrupted — goal paused (run /goal resume to continue)") + host.Goal().SetStatus(agenttool.GoalPaused) + } else { + fmt.Fprintf(out, "error: %v\n", err) + } + } +} + +// printGoalOutcome prints a one-line result banner after a goal run settles, +// keyed on the terminal status the run reached. +func printGoalOutcome(out io.Writer, snap agenttool.GoalSnapshot) { + color := ui.Enabled() + switch snap.Status { + case agenttool.GoalComplete: + fmt.Fprintf(out, "%s goal complete: %s\n", ui.Colorize(color, ui.Green, "✓"), snap.Summary) + case agenttool.GoalBlocked: + fmt.Fprintf(out, "%s goal blocked: %s\n", ui.Colorize(color, ui.Red, "⚠"), snap.BlockReason) + case agenttool.GoalBudgetLimited: + fmt.Fprintf(out, "%s goal paused — token budget reached (%d / %d). Run /goal resume to continue.\n", + ui.Colorize(color, ui.Yellow, "⏸"), snap.TokensUsed, snap.TokenBudget) + case agenttool.GoalPaused: + fmt.Fprintf(out, "%s goal paused after %d turns. Run /goal resume to continue, or /goal clear to drop it.\n", + ui.Colorize(color, ui.Yellow, "⏸"), snap.Iterations) + } +} + +// printGoalStatus prints a summary of the current goal state for a bare /goal. +func printGoalStatus(out io.Writer, goal *agenttool.GoalState) { + snap := goal.Snapshot() + if snap.Status == agenttool.GoalIdle { + fmt.Fprintln(out, "no goal set — run /goal to start one") + return + } + fmt.Fprintf(out, "goal: %s\n", snap.Objective) + fmt.Fprintf(out, "status: %s\n", snap.Status) + fmt.Fprintf(out, "iterations: %d\n", snap.Iterations) + if snap.TokenBudget > 0 { + fmt.Fprintf(out, "tokens: %d / %d\n", snap.TokensUsed, snap.TokenBudget) + } else { + fmt.Fprintf(out, "tokens: %d (no budget)\n", snap.TokensUsed) + } + if !snap.StartedAt.IsZero() { + fmt.Fprintf(out, "elapsed: %s\n", time.Since(snap.StartedAt).Round(time.Second)) + } + if snap.Summary != "" { + fmt.Fprintf(out, "summary: %s\n", snap.Summary) + } + if snap.BlockReason != "" { + fmt.Fprintf(out, "blocked: %s\n", snap.BlockReason) + } +} diff --git a/pigo/internal/cli/goal/goal_test.go b/pigo/internal/cli/goal/goal_test.go new file mode 100644 index 0000000..78c0049 --- /dev/null +++ b/pigo/internal/cli/goal/goal_test.go @@ -0,0 +1,141 @@ +// Tests for the /goal command's pure helpers: goalFollowUpDecision's branch +// logic (terminal states, token budget, turn cap, no-progress guard), +// parseGoalObjective / parseTokenBudget flag parsing, and goalTurnActivity's +// token/tool accounting. +package goal + +import ( + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" +) + +func TestGoalFollowUpDecision(t *testing.T) { + tests := []struct { + name string + snap agenttool.GoalSnapshot + wantCont bool + wantTerm agenttool.GoalStatus + }{ + { + name: "active continues", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalActive, Iterations: 1}, + wantCont: true, + wantTerm: agenttool.GoalIdle, + }, + { + name: "complete stops", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalComplete}, + wantCont: false, + wantTerm: agenttool.GoalIdle, + }, + { + name: "blocked stops", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalBlocked}, + wantCont: false, + wantTerm: agenttool.GoalIdle, + }, + { + name: "budget exhausted", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalActive, TokenBudget: 100, TokensUsed: 100}, + wantCont: false, + wantTerm: agenttool.GoalBudgetLimited, + }, + { + name: "budget under limit continues", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalActive, TokenBudget: 100, TokensUsed: 50}, + wantCont: true, + wantTerm: agenttool.GoalIdle, + }, + { + name: "turn cap", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalActive, Iterations: goalMaxAutomaticTurns}, + wantCont: false, + wantTerm: agenttool.GoalPaused, + }, + { + name: "no progress guard", + snap: agenttool.GoalSnapshot{Status: agenttool.GoalActive, NoProgress: goalMaxNoProgress}, + wantCont: false, + wantTerm: agenttool.GoalPaused, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cont, term := goalFollowUpDecision(tt.snap) + if cont != tt.wantCont { + t.Errorf("cont = %v, want %v", cont, tt.wantCont) + } + if term != tt.wantTerm { + t.Errorf("terminal = %q, want %q", term, tt.wantTerm) + } + }) + } +} + +func TestParseGoalObjective(t *testing.T) { + tests := []struct { + name string + args string + wantObj string + wantBud int + wantErr bool + }{ + {name: "no flag", args: "create hello.txt", wantObj: "create hello.txt", wantBud: 0}, + {name: "tokens k", args: "--tokens 100k build the thing", wantObj: "build the thing", wantBud: 100000}, + {name: "tokens m", args: "--tokens 1m do it", wantObj: "do it", wantBud: 1000000}, + {name: "tokens bare", args: "--tokens 5000 go", wantObj: "go", wantBud: 5000}, + {name: "flag only no objective", args: "--tokens 50k", wantObj: "", wantBud: 50000}, + {name: "bad value", args: "--tokens abc do it", wantErr: true}, + {name: "zero value", args: "--tokens 0 do it", wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + obj, bud, err := parseGoalObjective(tt.args) + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got obj=%q bud=%d", obj, bud) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if obj != tt.wantObj { + t.Errorf("objective = %q, want %q", obj, tt.wantObj) + } + if bud != tt.wantBud { + t.Errorf("budget = %d, want %d", bud, tt.wantBud) + } + }) + } +} + +func TestGoalTurnActivity(t *testing.T) { + tail := []agentcore.AgentMessage{ + agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Usage: &agentcore.Usage{OutputTokens: 40}, + }, + agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Usage: &agentcore.Usage{OutputTokens: 60}, + }, + } + tokens, hadTool := goalTurnActivity(tail) + if tokens != 100 { + t.Errorf("tokens = %d, want 100", tokens) + } + if hadTool { + t.Error("hadTool = true, want false (no tool calls or results)") + } + + withTool := []agentcore.AgentMessage{ + agentcore.ToolResultMessage{ToolName: "bash"}, + } + _, hadTool = goalTurnActivity(withTool) + if !hadTool { + t.Error("hadTool = false, want true (tool result present)") + } +} diff --git a/pigo/internal/cli/headless/headless.go b/pigo/internal/cli/headless/headless.go new file mode 100644 index 0000000..9797e3e --- /dev/null +++ b/pigo/internal/cli/headless/headless.go @@ -0,0 +1,205 @@ +// This file is the headless run driver: the print / stream-json run path +// (US-020) extracted from the CLI dispatch seam (#363). dispatch resolves the +// output mode and the run environment, then hands off to Run, which wires the +// session, prompt, thinking level, and provider credentials into a +// runtime.HeadlessConfig and executes one run. Plugin slash commands and output +// mode parsing live here because they are specific to the headless path. +package headless + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" +) + +// RunParams carries the resolved inputs for one headless run. Mode and Env are +// resolved by the caller (dispatch) — Mode via ParseOutputMode, Env via +// run.SetupEnv — so their distinct exit codes stay at the call site; Run owns +// the rest of the run lifecycle. +type RunParams struct { + Mode runtime.HeadlessMode + Env run.Env + Prompt string + Model string + APIKey string + ThinkingLevel string + ResumeID string +} + +// Run executes one headless run over p.Prompt, writing agent output to out and +// diagnostics to errOut, and returns a process exit code (0 = success). The run +// is backed by a session so its id appears in the first stream-json event and it +// can be resumed with --resume/--continue; a resumed session seeds its prior +// messages ahead of the new prompt. +func Run(ctx context.Context, p RunParams, out, errOut io.Writer) int { + env := p.Env + // Best-effort plugin slash-command support in headless mode: if the prompt is + // a "/cmd ..." naming a plugin command, invoke it, print its notifications to + // errOut, and use the returned prompt for this run (appending the raw args if + // the command produced no prompt). Headless has no turn injection, so + // appending the returned prompt is the accepted behavior. A non-plugin prompt + // or unknown command is left untouched. + headlessPrompt := resolveHeadlessPluginCommand(p.Prompt, env.Plugins, errOut) + promptContent, err := ui.BuildUserContent(headlessPrompt) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + + // Back the headless run with a session so its id appears in the first + // stream-json event and the run can be resumed with --resume/--continue, + // matching the interactive REPL and pi/Claude Code. A resumed session seeds + // its prior messages ahead of the new prompt. + priorMsgs, hs, err := openHeadlessSession(p.ResumeID, p.Model, env.ProviderName, env.SysPrompt) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + messages := append(priorMsgs, agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: promptContent}) + agentCtx := &agentcore.AgentContext{ + SystemPrompt: hs.header.SystemPrompt, + Messages: messages, + Tools: env.Tools, + } + + // Resolve the effective reasoning-effort level through the layered config + // chain (default < global < project < env < --thinking-level flag). + thinking, err := run.ResolveThinkingLevel(p.ThinkingLevel) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 2 + } + + // Resolve the API key by provider name from the environment (never logged). + // An explicit --api-key overrides env/config for the resolved provider. + creds := provider.NewCredentialStore(nil) + creds.SetOverride(env.ProviderName, p.APIKey) + runCfg := run.NewConfig(p.Model, env.ProviderName, thinking, env.Provider, creds, run.ToolRegistry(env.Tools), run.TodoReminders(env.Tools)) + runCfg.SessionID = hs.header.ID + // Route auto-compaction checkpoints to the shared memory root so a rebuild can + // recover the pre-watermark prefix (no-op when memory is disabled → empty root). + runCfg.MemoryRoot = run.MemoryRootFromTools(env.Tools) + + // Wire hooks uniformly with every other driver (#425): resolve the trust-gated + // hook set, install the tool-execution + Stop seams, dispatch SessionStart, and + // chain the SessionEnd/PreCompact observer onto the plugin event notifier. A + // malformed hook layer is a config error (exit 2), matching thinking-level. + source := "startup" + if p.ResumeID != "" { + source = "resume" + } + set, herr := run.ResolveHookSet(env.Cwd, run.Trusted(env.Cwd)) + if herr != nil { + fmt.Fprintf(errOut, "pigo: %v\n", herr) + return 2 + } + hookDeps := run.HookDeps{SessionID: hs.header.ID, ProjectDir: env.Cwd, WarnLog: errOut} + // Deliver agent lifecycle events to any subscribed plugin (US-017, #133). + // NewEventNotifier returns nil when no plugin subscribes, so the base handler + // stays nil in the common no-plugin case. + var baseOnEvent func(agentcore.AgentEvent) + if n := plugin.NewEventNotifier(env.Plugins, errOut); n != nil { + baseOnEvent = n.Handle + } + d, onEvent := run.InstallDriverHooks(ctx, &runCfg, set, hookDeps, source, baseOnEvent) + // UserPromptSubmit runs before the prompt is handed to the loop: a block aborts + // the headless run non-zero; additionalContext is injected into this run only. + if d != nil { + if block, reason := run.DispatchUserPromptSubmit(ctx, d, &runCfg, hookDeps, headlessPrompt); block { + fmt.Fprintf(errOut, "pigo: prompt blocked by hook: %s\n", reason) + return 1 + } + } + + cfg := runtime.HeadlessConfig{ + Mode: p.Mode, + Out: out, + Run: runCfg, + } + cfg.OnEvent = onEvent + runErr := runtime.RunHeadless(ctx, agentCtx, cfg) + // Persist the run's messages regardless of run outcome so a partial run is + // still resumable; a persistence failure is reported but does not mask a run + // error. + if perr := hs.persist(agentCtx); perr != nil { + fmt.Fprintf(errOut, "pigo: warning: could not persist session %s: %v\n", hs.header.ID, perr) + } + if runErr != nil { + fmt.Fprintf(errOut, "pigo: %v\n", runErr) + return 1 + } + return 0 +} + +// resolveHeadlessPluginCommand gives the headless / print path best-effort +// support for plugin slash commands. When prompt is a "/cmd ..." naming a +// plugin command (from mgr.Commands()), it invokes the command, prints each +// returned notification to notifyOut, and returns the command's returned Prompt +// as the run's prompt. If the command returns no prompt, the raw argument text +// is used instead (so a bare "/cmd" with only notifications still runs +// something sensible rather than an empty prompt). Any other input — a +// non-command, an unknown command, or a call error — leaves prompt unchanged so +// the normal headless run proceeds. mgr may be nil (no plugins). +// +// Headless has no turn-injection loop, so "inject the returned prompt" degrades +// to "use the returned prompt for this run", which the acceptance criteria +// permit. +func resolveHeadlessPluginCommand(prompt string, mgr *plugin.Manager, notifyOut io.Writer) string { + if mgr == nil || !strings.HasPrefix(strings.TrimLeft(prompt, " \t"), "/") { + return prompt + } + trimmed := strings.TrimLeft(prompt, " \t")[1:] + name := trimmed + args := "" + if i := strings.IndexAny(trimmed, " \t"); i >= 0 { + name = trimmed[:i] + args = strings.TrimSpace(trimmed[i+1:]) + } + for _, pc := range mgr.Commands() { + if pc.Spec.Name != name { + continue + } + // Encode the raw arg text as a JSON string (never null), matching the + // host's CommandCallParams.Args contract. + raw, _ := json.Marshal(args) + res, err := pc.Plugin.CallCommand(context.Background(), name, json.RawMessage(raw)) + if err != nil { + fmt.Fprintf(notifyOut, "pigo: plugin command %q failed: %v\n", name, err) + return prompt + } + for _, n := range res.Notifications { + if n.Type != "" { + fmt.Fprintf(notifyOut, "[%s] %s\n", n.Type, n.Message) + } else { + fmt.Fprintln(notifyOut, n.Message) + } + } + if res.Prompt != "" { + return res.Prompt + } + return args + } + return prompt +} + +// ParseOutputMode maps the --output-format flag onto a HeadlessMode, erroring on +// an unknown value. +func ParseOutputMode(outputFmt string) (runtime.HeadlessMode, error) { + switch outputFmt { + case "text", "": + return runtime.PrintMode, nil + case "stream-json": + return runtime.StreamJSONMode, nil + default: + return 0, fmt.Errorf("unknown --output-format %q (want text|stream-json)", outputFmt) + } +} diff --git a/pigo/internal/cli/headless/headless_test.go b/pigo/internal/cli/headless/headless_test.go new file mode 100644 index 0000000..588b617 --- /dev/null +++ b/pigo/internal/cli/headless/headless_test.go @@ -0,0 +1,41 @@ +package headless + +// Tests for the headless run driver's flag parsing. The run lifecycle (Run) is +// exercised via session/subagent tests here and provider-backed tests in +// internal/runtime; ParseOutputMode is pure and pinned directly. + +import ( + "testing" + + "github.com/smallnest/pigo/internal/runtime" +) + +// TestParseOutputMode covers the three accepted spellings and one rejection, +// pinning the flag contract the headless driver depends on. +func TestParseOutputMode(t *testing.T) { + cases := []struct { + in string + want runtime.HeadlessMode + wantErr bool + }{ + {"text", runtime.PrintMode, false}, + {"", runtime.PrintMode, false}, + {"stream-json", runtime.StreamJSONMode, false}, + {"yaml", 0, true}, + } + for _, c := range cases { + got, err := ParseOutputMode(c.in) + if c.wantErr { + if err == nil { + t.Errorf("ParseOutputMode(%q): want error, got nil", c.in) + } + continue + } + if err != nil { + t.Errorf("ParseOutputMode(%q): unexpected error %v", c.in, err) + } + if got != c.want { + t.Errorf("ParseOutputMode(%q) = %v, want %v", c.in, got, c.want) + } + } +} diff --git a/pigo/internal/cli/headless/session.go b/pigo/internal/cli/headless/session.go new file mode 100644 index 0000000..05964e7 --- /dev/null +++ b/pigo/internal/cli/headless/session.go @@ -0,0 +1,185 @@ +// Package headless drives pigo's non-interactive run paths: the print / +// stream-json headless run, the session listing/resume helpers, and the +// process-isolated sub-agent JSON-RPC server (--subagent-rpc). +// +// This file gives headless / stream-json runs the same session persistence and +// resume the interactive REPL has (cmd/pigo/interactive.go). Before this, a +// headless run built an in-memory AgentContext and threw it away on exit, so +// `--output-format stream-json` emitted no session id and `--resume`/`--continue` +// only worked in the REPL. +// +// Now a headless run is backed by a session file: resuming seeds the context +// from a prior session (and re-anchors the branch leaf), a fresh run creates a +// new session, and in both cases the run's newly produced messages are appended +// after it completes. The session id is threaded into the run so it appears in +// the first stream-json event (mirrors pi/Claude Code) and can be passed back via +// --resume to continue the run. +package headless + +import ( + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/session" +) + +// SessionStore returns the session store rooted at ~/.pigo/sessions (or under +// PIGO_HOME when set), creating the directory on first use. It is shared by the +// headless run path and the interactive REPL. +func SessionStore() (*session.Store, error) { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return nil, fmt.Errorf("resolve home dir: %w", err) + } + dir = filepath.Join(home, ".pigo") + } + return session.NewStore(filepath.Join(dir, "sessions")) +} + +// PrintSessions prints the stored sessions, most-recent first, to out. +func PrintSessions(out io.Writer) error { + store, err := SessionStore() + if err != nil { + return err + } + headers, err := store.List() + if err != nil { + return err + } + if len(headers) == 0 { + fmt.Fprintln(out, "no sessions") + return nil + } + for _, h := range headers { + fmt.Fprintf(out, "%s\t%s\t%s\n", h.ID, h.UpdatedAt.Local().Format("2006-01-02 15:04"), h.Model) + } + return nil +} + +// MostRecentSessionID returns the id of the most recently updated session, or +// "" if there are none. +func MostRecentSessionID() (string, error) { + store, err := SessionStore() + if err != nil { + return "", err + } + headers, err := store.List() + if err != nil { + return "", err + } + if len(headers) == 0 { + return "", nil + } + return headers[0].ID, nil +} + +// headlessSession is the session state backing one headless run: the store, the +// header (whose ID is the session id emitted and used for resume), and the +// branch-tracking cursor (curLeaf/persisted) so the run's messages append as a +// branch descending from the resumed leaf rather than flattening the tree. +type headlessSession struct { + store *session.Store + header session.SessionHeader + curLeaf string // active leaf id to descend from; "" for a fresh session + // persisted is the number of agentCtx.Messages already on disk before the + // run; persist appends only Messages[persisted:] as a new branch. + persisted int + // model/provider are the model and provider the run actually used, refreshed + // onto the header before persisting so a resumed run does not write back the + // original session's stale values (matching the REPL, repl.go persistTurn). + model string + provider string +} + +// openHeadlessSession resolves the session backing a headless run: it resumes an +// existing session when resumeID is set (seeding priorMsgs and re-anchoring the +// branch leaf) or creates a fresh session header otherwise. It returns the prior +// messages to seed into the context ahead of the new prompt, plus the session +// state used to persist the run afterward. +func openHeadlessSession(resumeID, model, providerName, sysPrompt string) (agentcore.MessageList, headlessSession, error) { + store, err := SessionStore() + if err != nil { + return nil, headlessSession{}, err + } + now := time.Now().UTC() + + if resumeID != "" { + h, entries, err := store.LoadEntries(resumeID) + if err != nil { + return nil, headlessSession{}, err + } + msgs := make(agentcore.MessageList, len(entries)) + for i, e := range entries { + msgs[i] = e.Message + } + curLeaf := "" + if len(entries) > 0 { + curLeaf = entries[len(entries)-1].ID + } + // A resumed header keeps its own SystemPrompt when present so the run is + // faithful to the original session. + if h.SystemPrompt == "" { + h.SystemPrompt = sysPrompt + } + return msgs, headlessSession{store: store, header: h, curLeaf: curLeaf, persisted: len(msgs), model: model, provider: providerName}, nil + } + + header := session.SessionHeader{ + ID: session.NewID(now), + CreatedAt: now, + UpdatedAt: now, + Model: model, + Provider: providerName, + SystemPrompt: sysPrompt, + Cwd: headlessCwd(), + } + return nil, headlessSession{store: store, header: header, curLeaf: "", persisted: 0, model: model, provider: providerName}, nil +} + +// headlessCwd returns the absolute working directory the run executes in, used +// to attribute the session to a project (SessionHeader.Cwd → project id) so a +// later /dream pass can distill this session under the right project scope. An +// unresolvable cwd yields "" (the session stays unattributed) rather than +// aborting the run. +func headlessCwd() string { + wd, err := os.Getwd() + if err != nil { + return "" + } + return wd +} + +// persist appends the messages produced during the run — everything in +// agentCtx.Messages past what was already on disk — as a branch descending from +// the resumed leaf, matching how the REPL grows a session tree (AppendBranch). +// It is a no-op when the run produced nothing new. Errors are returned for the +// caller to surface; the run's output has already been emitted regardless. +func (hs *headlessSession) persist(agentCtx *agentcore.AgentContext) error { + // Compaction can rebuild agentCtx.Messages to fewer entries than were on disk + // before the run (loop.go maybeAutoCompact replaces the slice). Clamp the + // cursor so the tail slice stays in bounds; when the context shrank there is + // nothing new to append past what compaction kept. + if hs.persisted > len(agentCtx.Messages) { + hs.persisted = len(agentCtx.Messages) + } + tail := agentCtx.Messages[hs.persisted:] + if len(tail) == 0 { + return nil + } + // Refresh the header with the model/provider the run actually used so a + // resumed session's metadata is not written back stale (matching the REPL). + hs.header.Model = hs.model + hs.header.Provider = hs.provider + hs.header.UpdatedAt = time.Now().UTC() + if _, err := hs.store.AppendBranch(hs.header, hs.curLeaf, tail); err != nil { + return err + } + hs.persisted = len(agentCtx.Messages) + return nil +} diff --git a/pigo/internal/cli/headless/session_test.go b/pigo/internal/cli/headless/session_test.go new file mode 100644 index 0000000..f1144d6 --- /dev/null +++ b/pigo/internal/cli/headless/session_test.go @@ -0,0 +1,156 @@ +package headless + +// Tests for headless session persistence and resume (session id in stream-json +// + --resume for headless runs). openHeadlessSession/persist are exercised +// directly against an isolated PIGO_HOME so a headless run's session round-trips +// without spawning a provider. + +import ( + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func textUser(s string) agentcore.UserMessage { + return agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent(s)}} +} + +func textAssistant(s string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent(s)}} +} + +// TestOpenHeadlessSessionFresh verifies a fresh headless session gets a new id +// and empty prior messages, and that persist writes the run's messages so they +// can be resumed. +func TestOpenHeadlessSessionFresh(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + + prior, hs, err := openHeadlessSession("", "faux-model", "faux", "sys prompt") + if err != nil { + t.Fatalf("openHeadlessSession fresh: %v", err) + } + if len(prior) != 0 { + t.Errorf("fresh session must have no prior messages, got %d", len(prior)) + } + if hs.header.ID == "" { + t.Fatal("fresh session must have a non-empty id") + } + if hs.header.SystemPrompt != "sys prompt" { + t.Errorf("header SystemPrompt = %q, want the passed prompt", hs.header.SystemPrompt) + } + + // Simulate a completed run: prompt + assistant reply. + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("1+1=?"), textAssistant("2")}} + if err := hs.persist(agentCtx); err != nil { + t.Fatalf("persist: %v", err) + } + + // The session must now be loadable with both messages. + _, msgs, err := hs.store.Load(hs.header.ID) + if err != nil { + t.Fatalf("Load persisted session: %v", err) + } + if len(msgs) != 2 { + t.Fatalf("persisted session has %d messages, want 2", len(msgs)) + } +} + +// TestOpenHeadlessSessionResume verifies that resuming seeds the prior messages +// and that a subsequent run appends only the new tail as a branch, so the +// session grows rather than being rewritten. +func TestOpenHeadlessSessionResume(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + + // First run: create and persist a session. + _, hs1, err := openHeadlessSession("", "faux-model", "faux", "sys") + if err != nil { + t.Fatalf("first openHeadlessSession: %v", err) + } + ctx1 := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("first"), textAssistant("reply1")}} + if err := hs1.persist(ctx1); err != nil { + t.Fatalf("first persist: %v", err) + } + sessID := hs1.header.ID + + // Second run: resume the session id. + prior, hs2, err := openHeadlessSession(sessID, "faux-model", "faux", "sys") + if err != nil { + t.Fatalf("resume openHeadlessSession: %v", err) + } + if len(prior) != 2 { + t.Fatalf("resume must seed %d prior messages, got %d", 2, len(prior)) + } + if hs2.header.ID != sessID { + t.Errorf("resumed session id = %q, want %q", hs2.header.ID, sessID) + } + if hs2.persisted != 2 { + t.Errorf("resumed persisted cursor = %d, want 2", hs2.persisted) + } + + // A second turn appends its new tail. + ctx2 := &agentcore.AgentContext{Messages: append(prior, textUser("second"), textAssistant("reply2"))} + if err := hs2.persist(ctx2); err != nil { + t.Fatalf("second persist: %v", err) + } + _, msgs, err := hs2.store.Load(sessID) + if err != nil { + t.Fatalf("Load after second turn: %v", err) + } + if len(msgs) != 4 { + t.Fatalf("session after two turns has %d messages, want 4", len(msgs)) + } +} + +// TestHeadlessPersistNoop verifies persist is a no-op (no error, no growth) when +// the run produced nothing new past what was already persisted. +func TestHeadlessPersistNoop(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + _, hs, err := openHeadlessSession("", "m", "p", "s") + if err != nil { + t.Fatalf("openHeadlessSession: %v", err) + } + ctx := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("x"), textAssistant("y")}} + if err := hs.persist(ctx); err != nil { + t.Fatalf("first persist: %v", err) + } + // Persisting again with no new messages must not error and must not duplicate. + if err := hs.persist(ctx); err != nil { + t.Fatalf("noop persist: %v", err) + } + _, msgs, err := hs.store.Load(hs.header.ID) + if err != nil { + t.Fatalf("Load: %v", err) + } + if len(msgs) != 2 { + t.Fatalf("noop persist changed message count to %d, want 2", len(msgs)) + } +} + +// TestHeadlessPersistCompactionShrink verifies persist tolerates the context +// being rebuilt to fewer messages than were on disk before the run (mid-run +// compaction replaces agentCtx.Messages). The persisted cursor is clamped so +// the tail slice stays in bounds rather than panicking. +func TestHeadlessPersistCompactionShrink(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + _, hs, err := openHeadlessSession("", "m", "p", "s") + if err != nil { + t.Fatalf("openHeadlessSession: %v", err) + } + // Persist four messages, advancing the cursor to 4. + ctx := &agentcore.AgentContext{Messages: agentcore.MessageList{textUser("a"), textAssistant("b"), textUser("c"), textAssistant("d")}} + if err := hs.persist(ctx); err != nil { + t.Fatalf("first persist: %v", err) + } + if hs.persisted != 4 { + t.Fatalf("cursor = %d, want 4", hs.persisted) + } + // Simulate compaction: the context is rebuilt to fewer messages than the + // cursor. persist must not panic on the out-of-range slice. + ctx.Messages = agentcore.MessageList{textAssistant("summary"), textUser("e")} + if err := hs.persist(ctx); err != nil { + t.Fatalf("persist after compaction shrink: %v", err) + } + if hs.persisted != 2 { + t.Errorf("cursor after clamp = %d, want 2", hs.persisted) + } +} diff --git a/pigo/internal/cli/headless/subagent_rpc.go b/pigo/internal/cli/headless/subagent_rpc.go new file mode 100644 index 0000000..9b63989 --- /dev/null +++ b/pigo/internal/cli/headless/subagent_rpc.go @@ -0,0 +1,161 @@ +// This file implements the subprocess side of process-isolated sub-agents +// (US-019, #135). Invoked as `pigo --subagent-rpc`, pigo speaks JSON-RPC 2.0 +// over stdio: for each "subagent/run" request on stdin it runs a child agent +// loop and writes the result (or an error) to stdout, exiting when stdin +// closes. The parent (SubAgentTool in process mode, internal/runtime) drives it +// via internal/jsonrpc. +// +// It reuses internal/jsonrpc's message types (Request/Response/ID/Version) for +// (de)serialization so the wire format matches the client exactly; the agent +// execution itself is runtime.RunSubAgentOnce. +package headless + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "os" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/jsonrpc" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" +) + +// RunSubAgentRPC is the `pigo --subagent-rpc` entry point. It reads +// newline-delimited JSON-RPC requests from in, runs each sub-agent request, and +// writes one response per request to out. It returns 0 (success) when stdin +// closes; a per-request failure is an RPC error response, not a non-zero exit, +// so the parent can distinguish "the child answered with an error" from "the +// child crashed" (the latter is detected by the parent's transport when stdout +// closes without a response). +func RunSubAgentRPC(ctx context.Context, in io.Reader, out, errOut io.Writer) int { + scanner := bufio.NewScanner(in) + // A sub-agent prompt can be large; match the jsonrpc client's line cap. + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + enc := json.NewEncoder(out) + for scanner.Scan() { + line := bytes.TrimSpace(scanner.Bytes()) + if len(line) == 0 { + continue + } + var req jsonrpc.Request + if err := json.Unmarshal(line, &req); err != nil { + // A parse error carries no id, so the response id is null. The + // jsonrpc client drops responses with a null id (it cannot correlate + // them), so this is only observable on the child's stderr; in + // practice the parent always sends well-formed requests. + writeSubAgentError(enc, nil, -32700, "parse error: "+err.Error()) + continue + } + handleSubAgentRequest(ctx, enc, &req) + } + // A scanner error (e.g. a request line exceeding the 16 MiB cap) ends the + // stream abnormally: surface it on stderr and exit non-zero so the parent's + // transport sees a diagnostic rather than a silent clean exit. + if err := scanner.Err(); err != nil { + fmt.Fprintf(errOut, "pigo: subagent-rpc stdin: %v\n", err) + return 1 + } + return 0 +} + +// handleSubAgentRequest dispatches one JSON-RPC request to the sub-agent runner +// and writes the response. Unknown methods, bad params, provider-resolution +// failures, and failed child runs are all RPC errors so the parent surfaces them +// as tool errors; only a successful run yields a result with the child's text. +func handleSubAgentRequest(ctx context.Context, enc *json.Encoder, req *jsonrpc.Request) { + if req.Method != runtime.SubAgentRPCMethod { + writeSubAgentError(enc, req.ID, -32601, "method not found: "+req.Method) + return + } + var params runtime.SubAgentRunParams + if len(req.Params) > 0 { + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + writeSubAgentError(enc, req.ID, -32602, "invalid params: "+err.Error()) + return + } + } + if params.Prompt == "" || params.Model == "" { + writeSubAgentError(enc, req.ID, -32602, "invalid params: prompt and model are required") + return + } + // Resolve the provider the same way the CLI does, so the subprocess targets + // the same gateway the parent's NewRunConfig encoded. Credentials come from + // the inherited environment (the parent's env vars). + prov, providerName, err := provider.ResolveProvider(params.Model, params.BaseURL, params.Protocol, "", os.Getenv) + if err != nil { + writeSubAgentError(enc, req.ID, -32603, "resolve provider: "+err.Error()) + return + } + cwd, _ := os.Getwd() + tools := filterBuiltinTools(run.BuiltinTools(cwd, false), params.Tools) + reg := run.ToolRegistry(tools) + creds := provider.NewCredentialStore(nil) // env-resolved + runCfg := runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: params.Model, + Provider: providerName, + Stream: provider.StreamFnFromProvider(prov), + GetAPIKey: creds.GetAPIKey, + }, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}}, + } + // Wire hooks uniformly with every other driver (#425): the child sub-agent runs + // its own PreToolUse/PostToolUse (and Stop) hooks from the trust-gated hook set + // rooted at its working directory. It has no backing session, so SessionID is + // empty (omitted from HookInput). A malformed hook layer disables hooks with a + // warning rather than failing the child run. + if set, herr := run.ResolveHookSet(cwd, run.Trusted(cwd)); herr != nil { + fmt.Fprintf(os.Stderr, "pigo: hooks disabled: %v\n", herr) + } else { + run.InstallHooks(&runCfg, set, run.HookDeps{ProjectDir: cwd, WarnLog: os.Stderr}) + } + text, err := runtime.RunSubAgentOnce(ctx, params.SystemPrompt, params.Prompt, tools, runCfg) + if err != nil { + // A failed child run is an RPC error so the parent's defaultProcessCall + // returns a Go error and executeProcess marks the tool result IsError, + // matching goroutine mode's "failed run -> tool error" behavior. + writeSubAgentError(enc, req.ID, -32000, err.Error()) + return + } + result, _ := json.Marshal(runtime.SubAgentRunResult{Text: text}) + _ = enc.Encode(jsonrpc.Response{JSONRPC: jsonrpc.Version, ID: req.ID, Result: result}) +} + +// writeSubAgentError writes a JSON-RPC error response with the given id (which +// may be nil for a parse error on an unidentifiable request) and code/message. +func writeSubAgentError(enc *json.Encoder, id *jsonrpc.ID, code int, msg string) { + _ = enc.Encode(jsonrpc.Response{ + JSONRPC: jsonrpc.Version, + ID: id, + Error: &jsonrpc.Error{Code: code, Message: msg}, + }) +} + +// filterBuiltinTools returns the subset of tools whose Name is in names. An +// empty names list keeps all tools. It lets a process-isolated sub-agent +// restrict the child to a subset (e.g. a read-only researcher), matching +// goroutine mode's Tools filtering, without serializing in-process tool objects +// across the process boundary. +func filterBuiltinTools(tools []agentcore.AgentTool, names []string) []agentcore.AgentTool { + if len(names) == 0 { + return tools + } + want := make(map[string]bool, len(names)) + for _, n := range names { + want[n] = true + } + var out []agentcore.AgentTool + for _, t := range tools { + if want[t.Name()] { + out = append(out, t) + } + } + return out +} diff --git a/pigo/internal/cli/headless/subagent_rpc_test.go b/pigo/internal/cli/headless/subagent_rpc_test.go new file mode 100644 index 0000000..a087c0e --- /dev/null +++ b/pigo/internal/cli/headless/subagent_rpc_test.go @@ -0,0 +1,120 @@ +package headless + +// Tests for the sub-agent RPC subprocess mode (US-019, #135): the pure +// filterBuiltinTools helper and the RunSubAgentRPC validation branches +// (method-not-found, invalid params, parse error) that return RPC errors before +// any provider is resolved. The happy-path transport is covered in +// internal/runtime via a compiled helper binary; RunSubAgentOnce (the agent +// core) is covered there with a faux provider. + +import ( + "bytes" + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/jsonrpc" + "github.com/smallnest/pigo/internal/runtime" +) + +// TestFilterBuiltinTools verifies the subprocess tool filter: an empty name +// list keeps all builtins, a subset keeps only the named tools, and unknown +// names are silently ignored. +func TestFilterBuiltinTools(t *testing.T) { + all := run.BuiltinTools(t.TempDir(), false) + if len(all) == 0 { + t.Fatal("BuiltinTools returned no tools") + } + namesOf := func(ts []agentcore.AgentTool) []string { + out := make([]string, len(ts)) + for i, tl := range ts { + out[i] = tl.Name() + } + return out + } + + if got := filterBuiltinTools(all, nil); len(got) != len(all) { + t.Errorf("nil names kept %d, want all %d", len(got), len(all)) + } + if got := filterBuiltinTools(all, []string{}); len(got) != len(all) { + t.Errorf("empty names kept %d, want all %d", len(got), len(all)) + } + + got := filterBuiltinTools(all, []string{"read", "grep"}) + if len(got) != 2 { + t.Fatalf("subset kept %d, want 2: %v", len(got), namesOf(got)) + } + gotNames := namesOf(got) + if gotNames[0] != "read" || gotNames[1] != "grep" { + t.Errorf("subset names = %v, want [read grep]", gotNames) + } + + // Unknown names are ignored, known ones kept. + got = filterBuiltinTools(all, []string{"read", "does-not-exist"}) + if len(got) != 1 || got[0].Name() != "read" { + t.Errorf("unknown-name filter kept %v, want [read]", namesOf(got)) + } +} + +// TestRunSubAgentRPCValidation verifies the subprocess returns JSON-RPC errors +// for malformed requests without reaching provider resolution: a parse error, +// an unknown method, and missing prompt/model params. These branches are the +// server's contract for bad input and need not involve a provider. +func TestRunSubAgentRPCValidation(t *testing.T) { + cases := []struct { + name string + line string + wantCode int + wantMsg string + }{ + {"parse error", "not-json", -32700, "parse error"}, + {"method not found", `{"jsonrpc":"2.0","id":1,"method":"other","params":{}}`, -32601, "method not found"}, + {"missing prompt", `{"jsonrpc":"2.0","id":2,"method":"` + runtime.SubAgentRPCMethod + `","params":{"model":"x"}}`, -32602, "prompt and model are required"}, + {"missing model", `{"jsonrpc":"2.0","id":3,"method":"` + runtime.SubAgentRPCMethod + `","params":{"prompt":"x"}}`, -32602, "prompt and model are required"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + var out, errOut bytes.Buffer + code := RunSubAgentRPC(context.Background(), strings.NewReader(c.line+"\n"), &out, &errOut) + if code != 0 { + t.Errorf("exit code = %d, want 0 (validation errors are RPC responses, not non-zero exits)", code) + } + line, _ := out.ReadString('\n') + if line == "" { + t.Fatal("no response written") + } + var resp jsonrpc.Response + if err := json.Unmarshal([]byte(line), &resp); err != nil { + t.Fatalf("unmarshal response %q: %v", line, err) + } + if resp.Error == nil { + t.Fatalf("response has no error: %s", line) + } + if resp.Error.Code != c.wantCode { + t.Errorf("error code = %d, want %d (msg=%q)", resp.Error.Code, c.wantCode, resp.Error.Message) + } + if !strings.Contains(resp.Error.Message, c.wantMsg) { + t.Errorf("error msg = %q, want it to contain %q", resp.Error.Message, c.wantMsg) + } + }) + } +} + +// TestRunSubAgentRPCScannerError verifies a stdin read error (a line exceeding +// the scanner cap) is surfaced on stderr and yields a non-zero exit, rather +// than a silent clean exit. +func TestRunSubAgentRPCScannerError(t *testing.T) { + // A line longer than the 16 MiB scanner cap triggers a scanner error. + huge := strings.Repeat("a", 17*1024*1024) + var out, errOut bytes.Buffer + code := RunSubAgentRPC(context.Background(), strings.NewReader(huge), &out, &errOut) + if code == 0 { + t.Error("exit code = 0 on scanner error, want non-zero") + } + if !strings.Contains(errOut.String(), "subagent-rpc stdin") { + t.Errorf("stderr = %q, want it to mention 'subagent-rpc stdin'", errOut.String()) + } +} diff --git a/pigo/internal/cli/host.go b/pigo/internal/cli/host.go new file mode 100644 index 0000000..b96b01b --- /dev/null +++ b/pigo/internal/cli/host.go @@ -0,0 +1,91 @@ +// This file defines the Host and Editor contracts that let the /goal, /btw, +// /status and REPL subpackages read and mutate a session's live state without +// importing the concrete replDeps aggregate that assembles it (see doc.go). The +// aggregate implements Host by exposing accessor and mutator methods over its +// fields; the line editor implements Editor. +package cli + +import ( + "bufio" + "errors" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +// ErrLineInterrupted is returned by an Editor.ReadLine when the user hits an +// idle Ctrl+C at the prompt (rather than during a run). Control commands that +// run their own follow-up loop (e.g. /btw) check for it with errors.Is to leave +// the loop cleanly. It lives here, beside the Editor contract, so a subpackage +// can recognize the interrupt without importing the concrete line editor. +var ErrLineInterrupted = errors.New("line input interrupted") + +// Host is the seam through which a control command (/goal, /btw, /status) and +// the REPL loop access the collaborators and mutable state of a session. It is +// satisfied by the replDeps aggregate assembled once per session. +// +// Getters return the shared collaborators (never copies, except where noted). +// The setters cover the fields a command may advance mid-session: the active +// session leaf, the persisted-message count, and the last /btw side thread. +type Host interface { + // Session collaborators. + Store() *session.Store + Header() session.SessionHeader + AgentCtx() *agentcore.AgentContext + Live() *LiveConfig + Registry() *agenttool.ToolRegistry + Reminders() *runtime.ReminderRegistry + Slash() *runtime.SlashRegistry + Creds() *provider.CredentialStore + Notifier() *plugin.EventNotifier + // NotifierHandle returns the plugin event-delivery callback for this + // session's runs, or nil when no plugin subscribed. + NotifierHandle() func(agentcore.AgentEvent) + Trust() *trust.Manager + Goal() *agenttool.GoalState + Telemetry() *TelemetryHolder + + // Dispatcher returns the session's hook dispatcher, or nil when no hooks are + // configured (FR-18). Side-runs (/goal, /btw) install the per-turn seams onto + // their own cfg via run.InstallSeams(cfg, host.Dispatcher(), host.HookDeps()). + Dispatcher() *hooks.Dispatcher + // HookDeps carries the session id / project dir stamped onto every HookInput. + HookDeps() run.HookDeps + + // Cwd is the directory pigo was launched in; it does not change during a + // session and gates side-effect tools. + Cwd() string + // Input is the shared buffered stdin reader used by both the main loop and + // the tool-call confirmation prompt. + Input() *bufio.Reader + // ConfirmMu serializes tool-call confirmation prompts across concurrent + // side-effect tool calls. + ConfirmMu() *sync.Mutex + + // Session-tree cursor and persistence bookkeeping. + CurLeaf() string + SetCurLeaf(id string) + Persisted() int + SetPersisted(n int) + + // Last /btw side thread from this process and its background base index. + LastBtw() *agentcore.AgentContext + SetLastBtw(ctx *agentcore.AgentContext) + LastBtwBase() int + SetLastBtwBase(n int) +} + +// Editor is the line-input contract a control command uses to read a follow-up +// line from the user (e.g. the /btw follow-up loop). It is satisfied by the +// REPL's line editor. +type Editor interface { + ReadLine(prompt string) (string, error) +} diff --git a/pigo/internal/cli/liveconfig.go b/pigo/internal/cli/liveconfig.go new file mode 100644 index 0000000..fa93532 --- /dev/null +++ b/pigo/internal/cli/liveconfig.go @@ -0,0 +1,38 @@ +// This file defines LiveConfig, the mutable run configuration a control command +// may change mid-session. It was moved verbatim from cmd/pigo (the former +// liveRunConfig) and exported so the run, repl, btw, status and goal +// subpackages can read and mutate it through the Host contract. The run closure +// reads it on every prompt, so a /model switch takes effect on the next turn. +// It carries no lock: it is read and written only on the REPL's single main +// goroutine (slash actions and the run are both invoked synchronously from the +// REPL loop, never concurrently). +package cli + +import ( + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// LiveConfig is the mutable run configuration a control command may change +// mid-session. +type LiveConfig struct { + Model string + ProviderName string + Provider provider.Provider + BaseURL string + Protocol string + // ThinkingLevel is the reasoning-effort level applied to each turn. It is + // seeded from the resolved config chain and read on every prompt. + ThinkingLevel agentcore.ThinkingLevel + // ContextWindow is the model's total context-token budget, used to gate + // automatic compaction. When 0 the window is unknown and auto-compaction is + // disabled; the REPL seeds it with a conservative default so long sessions + // still compact rather than overflow. + ContextWindow int +} + +// DefaultContextWindow is the fallback context-token budget used when a model's +// true window is unknown. It is deliberately large so auto-compaction only fires +// on genuinely long sessions (threshold = window - ReserveTokens), never on +// ordinary short exchanges. +const DefaultContextWindow = 128000 diff --git a/pigo/internal/cli/memstatus/memstatus.go b/pigo/internal/cli/memstatus/memstatus.go new file mode 100644 index 0000000..02ddf99 --- /dev/null +++ b/pigo/internal/cli/memstatus/memstatus.go @@ -0,0 +1,195 @@ +// This file implements the /memory slash command (US-011, FR-20, #484) that +// prints a colored report describing the persistent memory store and the +// infinite-context state: entry counts by scope, the current context window and +// auto-compaction trigger point, current context usage, and checkpoint status. +// +// Unlike /status (which reads everything through cli.Host), /memory takes its +// inputs explicitly so both the REPL and the TUI can call it with the live +// memory store, which the Host contract does not expose. +package memstatus + +import ( + "fmt" + "io" + "sort" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/memory" + "github.com/smallnest/pigo/internal/runtime" +) + +// scopeOrder is the display order for memory scopes. Scopes not listed here are +// appended afterwards in lexical order. +var scopeOrder = []memory.Scope{ + memory.ScopeGlobal, + memory.ScopeProjects, + memory.ScopeSessions, + memory.ScopeCC, +} + +// RunMemory prints a colored multi-section report about persistent memory and +// infinite context to out. store may be nil (persistent memory disabled), in +// which case the memory-store section reports that state and the checkpoint / +// context sections still render from msgs and the checkpoint file. +func RunMemory(out io.Writer, store *memory.Store, memoryRoot, sessionID string, msgs agentcore.MessageList, contextWindow int) { + color := ui.Enabled() + + fmt.Fprintln(out) + printStoreSection(out, color, store, memoryRoot) + fmt.Fprintln(out) + printContextSection(out, color, msgs, contextWindow) + fmt.Fprintln(out) + printCheckpointSection(out, color, memoryRoot, sessionID) +} + +// printStoreSection prints the persistent memory store section: whether memory +// is enabled, the root directory, and entry counts per scope. It reconciles the +// on-disk files into the index first so the counts are fresh. +func printStoreSection(out io.Writer, color bool, store *memory.Store, memoryRoot string) { + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "persistent memory:")) + if store == nil { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "status:"), + ui.Colorize(color, ui.Yellow, "disabled")) + return + } + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "status:"), + ui.Colorize(color, ui.Green, "enabled")) + if memoryRoot != "" { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "root:"), memoryRoot) + } + + if _, err := store.Reconcile(); err != nil { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "entries:"), + ui.Colorize(color, ui.Red, fmt.Sprintf("reconcile failed: %v", err))) + return + } + counts, err := store.CountByScope() + if err != nil { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "entries:"), + ui.Colorize(color, ui.Red, fmt.Sprintf("count failed: %v", err))) + return + } + + total := 0 + for _, n := range counts { + total += n + } + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "entries:"), total) + for _, scope := range orderedScopes(counts) { + fmt.Fprintf(out, " %s %d\n", + ui.Colorize(color, ui.Dim, string(scope)+":"), counts[scope]) + } +} + +// orderedScopes returns the scopes present in counts, listing the well-known +// scopes first (scopeOrder) and any others afterwards in lexical order. +func orderedScopes(counts map[memory.Scope]int) []memory.Scope { + seen := make(map[memory.Scope]bool, len(counts)) + var out []memory.Scope + for _, s := range scopeOrder { + if _, ok := counts[s]; ok { + out = append(out, s) + seen[s] = true + } + } + var extra []memory.Scope + for s := range counts { + if !seen[s] { + extra = append(extra, s) + } + } + sort.Slice(extra, func(i, j int) bool { return extra[i] < extra[j] }) + return append(out, extra...) +} + +// printContextSection prints the current context usage and the auto-compaction +// trigger point (window minus reserve), mirroring /status's context block so +// /memory is self-contained for infinite-context inspection. +func printContextSection(out io.Writer, color bool, msgs agentcore.MessageList, contextWindow int) { + tokens := compaction.EstimateContextTokens(msgs).Tokens + + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "context:")) + if contextWindow > 0 { + fmt.Fprintf(out, " %s %d / %d tokens\n", + ui.Colorize(color, ui.Dim, "current:"), tokens, contextWindow) + util := int(float64(tokens) / float64(contextWindow) * 100) + utilColor := ui.Green + if util >= 90 { + utilColor = ui.Red + } else if util >= 70 { + utilColor = ui.Yellow + } + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "utilization:"), + ui.Colorize(color, utilColor, fmt.Sprintf("%d%%", util))) + + reserve := compaction.DefaultCompactionSettings.ReserveTokens + threshold := contextWindow - reserve + remaining := threshold - tokens + if remaining < 0 { + fmt.Fprintf(out, " %s %s (window: %d, reserve: %d)\n", + ui.Colorize(color, ui.Dim, "compaction trigger:"), + ui.Colorize(color, ui.Red, fmt.Sprintf("%d tokens over trigger (%d)", -remaining, threshold)), + contextWindow, reserve) + } else { + fmt.Fprintf(out, " %s %s (window: %d, reserve: %d)\n", + ui.Colorize(color, ui.Dim, "compaction trigger:"), + ui.Colorize(color, ui.Green, fmt.Sprintf("%d tokens until trigger (%d)", remaining, threshold)), + contextWindow, reserve) + } + } else { + fmt.Fprintf(out, " %s %d tokens\n", ui.Colorize(color, ui.Dim, "current:"), tokens) + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "compaction trigger:"), + ui.Colorize(color, ui.Yellow, "auto-compaction disabled (unknown window)")) + } +} + +// printCheckpointSection prints the infinite-context checkpoint status for the +// session: whether a checkpoint.md exists and, if so, its watermark, covered +// message count, and creation time. +func printCheckpointSection(out io.Writer, color bool, memoryRoot, sessionID string) { + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "checkpoint:")) + if memoryRoot == "" || sessionID == "" { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "status:"), + ui.Colorize(color, ui.Yellow, "unavailable (no memory root or session id)")) + return + } + cp, ok, err := runtime.LoadCheckpoint(sessionID, memoryRoot) + if err != nil { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "status:"), + ui.Colorize(color, ui.Red, fmt.Sprintf("load failed: %v", err))) + return + } + if !ok { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "status:"), + ui.Colorize(color, ui.Dim, "none yet")) + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "path:"), + runtime.CheckpointPath(sessionID, memoryRoot)) + return + } + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "status:"), + ui.Colorize(color, ui.Green, "present")) + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "watermark:"), cp.Watermark) + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "covered messages:"), cp.CoveredMessages) + if !cp.CreatedAt.IsZero() { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "created:"), + cp.CreatedAt.Format("2006-01-02 15:04:05 MST")) + } + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "path:"), + runtime.CheckpointPath(sessionID, memoryRoot)) +} + diff --git a/pigo/internal/cli/memstatus/memstatus_test.go b/pigo/internal/cli/memstatus/memstatus_test.go new file mode 100644 index 0000000..d86d706 --- /dev/null +++ b/pigo/internal/cli/memstatus/memstatus_test.go @@ -0,0 +1,53 @@ +package memstatus + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/memory" +) + +// TestRunMemoryDisabled renders the disabled state when the store is nil and +// still prints the context and checkpoint sections. +func TestRunMemoryDisabled(t *testing.T) { + var buf bytes.Buffer + RunMemory(&buf, nil, "", "", nil, 0) + out := buf.String() + for _, want := range []string{"persistent memory:", "disabled", "context:", "checkpoint:"} { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } +} + +// TestRunMemoryEnabledCounts reconciles a store with entries and asserts the +// report shows enabled status and per-scope counts. +func TestRunMemoryEnabledCounts(t *testing.T) { + base := t.TempDir() + root := filepath.Join(base, "mem") + if err := os.MkdirAll(filepath.Join(root, "global", "reference"), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + if err := os.WriteFile(filepath.Join(root, "global", "reference", "a.md"), []byte("hello"), 0o644); err != nil { + t.Fatalf("write: %v", err) + } + st, err := memory.Open(filepath.Join(base, "index.db"), root, "") + if err != nil { + t.Fatalf("Open: %v", err) + } + defer st.Close() + + var buf bytes.Buffer + msgs := agentcore.MessageList{agentcore.UserMessage{Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}} + RunMemory(&buf, st, root, "sess-1", msgs, 200000) + out := buf.String() + for _, want := range []string{"enabled", "entries:", "global:", "context:", "checkpoint:", "none yet"} { + if !strings.Contains(out, want) { + t.Fatalf("output missing %q:\n%s", want, out) + } + } +} diff --git a/pigo/internal/cli/persist.go b/pigo/internal/cli/persist.go new file mode 100644 index 0000000..007e8bf --- /dev/null +++ b/pigo/internal/cli/persist.go @@ -0,0 +1,53 @@ +// This file holds PersistTurn, the session-tree persistence step shared by the +// REPL loop and the /goal autonomous loop. It reads and advances a session's +// mutable cursor state through the Host contract, so a command need not import +// the concrete replDeps aggregate to persist the tail of a turn. +package cli + +import ( + "fmt" + "io" + "time" +) + +// PersistTurn appends the messages produced since the last persist as a new +// branch descending from the host's current leaf, advancing the leaf and the +// persisted-message cursor. Growing the tree with AppendBranch (rather than a +// full linear Save) is what lets a later /tree leaf-switch fork the on-disk +// history instead of clobbering it. If nothing new was produced it is a no-op: +// rewriting the file would regenerate entry ids and flatten the tree. +func PersistTurn(out io.Writer, h Host) { + agentCtx := h.AgentCtx() + // Automatic compaction during a turn rewrites Messages into a summary + recent + // tail, shrinking the slice below the persisted cursor. The append-a-tail + // branch model no longer holds (the prefix changed and Messages[persisted:] + // would be out of range), so re-save the flattened context linearly and reset + // the branch cursor to the new leaf, mirroring the /compact handler. + if h.Persisted() > len(agentCtx.Messages) { + header := h.Header() + header.UpdatedAt = time.Now().UTC() + if err := h.Store().Save(header, agentCtx.Messages); err != nil { + fmt.Fprintf(out, "pigo: session save failed: %v\n", err) + return + } + h.SetPersisted(len(agentCtx.Messages)) + h.SetCurLeaf("") + if _, entries, err := h.Store().LoadEntries(header.ID); err == nil && len(entries) > 0 { + h.SetCurLeaf(entries[len(entries)-1].ID) + } + return + } + tail := agentCtx.Messages[h.Persisted():] + if len(tail) == 0 { + return + } + header := h.Header() + header.UpdatedAt = time.Now().UTC() + leaf, err := h.Store().AppendBranch(header, h.CurLeaf(), tail) + if err != nil { + fmt.Fprintf(out, "pigo: session save failed: %v\n", err) + return + } + h.SetCurLeaf(leaf) + h.SetPersisted(len(agentCtx.Messages)) +} diff --git a/pigo/internal/cli/pkgcmd/pkgcmd.go b/pigo/internal/cli/pkgcmd/pkgcmd.go new file mode 100644 index 0000000..53ee9c8 --- /dev/null +++ b/pigo/internal/cli/pkgcmd/pkgcmd.go @@ -0,0 +1,152 @@ +// Package pkgcmd wires pigo's package-management subcommands (#162, #163, #164) +// into the CLI: `pigo install|list|uninstall|update ...`. These are positional +// subcommands, distinct from the flag-driven agent modes, so main() peels them +// off before pflag parsing (the agent flags don't apply to package management). +// +// Each subcommand is a thin shell over internal/pkgmgr, which owns all the real +// work (fetch, classify, distribute, lockfile). This file only parses argv, +// resolves the lockfile path, calls pkgmgr, and prints human-readable results. +package pkgcmd + +import ( + "fmt" + "io" + + "github.com/smallnest/pigo/internal/pkgmgr" +) + +// Subcommands are the argv[1] values routed to Run. +var Subcommands = map[string]bool{ + "install": true, + "list": true, + "uninstall": true, + "update": true, +} + +// Run executes a package-management subcommand (cmd) with its arguments (args), +// writing output to out and errors to errOut. It returns a process exit code. +// install is #162; list/uninstall are #163; update is #164. +func Run(cmd string, args []string, out, errOut io.Writer) int { + lockPath := pkgmgr.DefaultLockfilePath() + switch cmd { + case "install": + return runInstall(args, lockPath, out, errOut) + case "list": + return runList(lockPath, out, errOut) + case "uninstall": + return runUninstall(args, lockPath, out, errOut) + case "update": + return runUpdate(args, lockPath, out, errOut) + default: + fmt.Fprintf(errOut, "pigo: %q is not yet implemented\n", cmd) + return 2 + } +} + +// runList handles `pigo list`: it prints one line per installed package +// (name, version, types, source), or a friendly notice when none are installed. +func runList(lockPath string, out, errOut io.Writer) int { + pkgs, err := pkgmgr.ListInstalled(lockPath) + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + if len(pkgs) == 0 { + fmt.Fprintln(out, "no packages installed") + return 0 + } + for _, p := range pkgs { + fmt.Fprintf(out, "%s\t%s\t%s\t%s\n", p.Name, p.Version, joinPkgTypes(p.Types), p.Source) + } + return 0 +} + +// runUninstall handles `pigo uninstall [more...]`. It removes each named +// package's files and lockfile entry, continuing past a failure so one bad name +// does not abort the rest; the exit code is non-zero if any uninstall failed. +func runUninstall(names []string, lockPath string, out, errOut io.Writer) int { + if len(names) == 0 { + fmt.Fprintln(errOut, "pigo: uninstall requires a package name, e.g. pigo uninstall pi-mcp-adapter") + return 2 + } + failed := false + for _, name := range names { + if err := pkgmgr.Uninstall(name, lockPath, out); err != nil { + fmt.Fprintf(errOut, "pigo: uninstall %s failed: %v\n", name, err) + failed = true + continue + } + fmt.Fprintf(out, "Uninstalled %s\n", name) + } + if failed { + return 1 + } + return 0 +} + +// runInstall handles `pigo install npm:[@version] [more...]`. It requires +// npm on PATH (checked once up front for a clear early failure) and installs +// each reference in turn, continuing past a failure so one bad package does not +// abort the rest; the exit code is non-zero if any install failed. +func runInstall(refs []string, lockPath string, out, errOut io.Writer) int { + if len(refs) == 0 { + fmt.Fprintln(errOut, "pigo: install requires a package reference, e.g. pigo install npm:pi-mcp-adapter") + return 2 + } + if err := pkgmgr.EnsureNPM(); err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + + failed := false + for _, ref := range refs { + res, err := pkgmgr.Install(ref, lockPath, out) + if err != nil { + fmt.Fprintf(errOut, "pigo: install %s failed: %v\n", ref, err) + failed = true + continue + } + fmt.Fprintf(out, "Installed %s@%s (%s) — %d file(s)\n", + res.Name, res.Version, joinPkgTypes(res.Types), len(res.Files)) + } + if failed { + return 1 + } + return 0 +} + +// runUpdate handles `pigo update [more...]`: it updates each named +// package, continuing past a failure so one bad package does not abort the rest. +// The exit code is non-zero if any update failed. A no-name `pigo update` is +// binary self-update, routed away by main() before pkgcmd is reached (US-003), +// so here an empty name list is a usage error rather than an update-all. +func runUpdate(names []string, lockPath string, out, errOut io.Writer) int { + if len(names) == 0 { + fmt.Fprintln(errOut, "pigo: update requires a package name, e.g. pigo update pi-mcp-adapter") + return 2 + } + failed := false + for _, name := range names { + if _, err := pkgmgr.Update(name, lockPath, out); err != nil { + fmt.Fprintf(errOut, "pigo: update %s failed: %v\n", name, err) + failed = true + continue + } + } + if failed { + return 1 + } + return 0 +} + +// joinPkgTypes renders package types as a comma-separated string. +func joinPkgTypes(types []pkgmgr.PackageType) string { + out := "" + for i, t := range types { + if i > 0 { + out += ", " + } + out += string(t) + } + return out +} diff --git a/pigo/internal/cli/prompts/presets_test.go b/pigo/internal/cli/prompts/presets_test.go new file mode 100644 index 0000000..f019c23 --- /dev/null +++ b/pigo/internal/cli/prompts/presets_test.go @@ -0,0 +1,33 @@ +package prompts + +// Tests for the /models preset listing. presetListing renders the curated +// catalog for the REPL /models command. Provider-resolution tests moved to +// internal/provider with the resolution logic itself (US-004, #361). + +import ( + "strings" + "testing" +) + +// TestPresetListingGroupsAndFilters verifies /models lists all providers by +// default and filters to one provider when given an argument. +func TestPresetListingGroupsAndFilters(t *testing.T) { + all := presetListing("") + for _, want := range []string{"openrouter", "nvidia", "ollama"} { + if !strings.Contains(all, want) { + t.Errorf("full listing missing provider %q:\n%s", want, all) + } + } + // Filter to nvidia only: openrouter must not appear. + nv := presetListing("nvidia") + if !strings.Contains(nv, "nvidia") { + t.Errorf("filtered listing missing nvidia:\n%s", nv) + } + if strings.Contains(nv, "openrouter") { + t.Errorf("nvidia filter must not include openrouter:\n%s", nv) + } + // Unknown filter yields a helpful message, not a crash. + if got := presetListing("bogus"); !strings.Contains(got, "no preset provider") { + t.Errorf("unknown filter = %q, want a not-found message", got) + } +} diff --git a/pigo/internal/cli/prompts/prompts_cli_test.go b/pigo/internal/cli/prompts/prompts_cli_test.go new file mode 100644 index 0000000..016fb12 --- /dev/null +++ b/pigo/internal/cli/prompts/prompts_cli_test.go @@ -0,0 +1,116 @@ +package prompts + +// Tests for --prompt-template (CLI tier) and --no-prompt-templates (US-008, +// #339): CLI paths load at the CLI tier, --no-prompt-templates suppresses all +// prompt-template discovery while leaving built-ins, and a global template +// overrides a same-named CLI one (global tier wins, CLI shadowed). + +import ( + "os" + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/testutil" + "github.com/smallnest/pigo/internal/runtime" +) + +func TestBuildSlashRegistryLoadsCLIPrompts(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + // file entry + filePath := filepath.Join(home, "single.md") + if err := os.WriteFile(filePath, []byte("Single: $ARGUMENTS"), 0o644); err != nil { + t.Fatal(err) + } + // dir entry + dirPath := filepath.Join(home, "clidir") + if err := os.MkdirAll(dirPath, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dirPath, "a.md"), []byte("A: $1"), 0o644); err != nil { + t.Fatal(err) + } + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{CLI: []string{filePath, dirPath}}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + out, err := reg.ResolveOutcome("/single hi") + if err != nil { + t.Fatalf("ResolveOutcome /single: %v", err) + } + if !out.Handled || out.Prompt != "Single: hi" { + t.Errorf("/single = handled=%v prompt=%q, want \"Single: hi\"", out.Handled, out.Prompt) + } + out2, err := reg.ResolveOutcome("/a x") + if err != nil { + t.Fatalf("ResolveOutcome /a: %v", err) + } + if !out2.Handled || out2.Prompt != "A: x" { + t.Errorf("/a = handled=%v prompt=%q, want \"A: x\"", out2.Handled, out2.Prompt) + } +} + +func TestBuildSlashRegistryNoPromptTemplatesDisables(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + // A global prompt that should NOT load under --no-prompt-templates. + testutil.WritePrompt(t, home, "prompts", "review.md", "Review: $ARGUMENTS") + // A CLI path that should also be ignored under --no-prompt-templates. + cliFile := filepath.Join(home, "cli.md") + if err := os.WriteFile(cliFile, []byte("CLI: $ARGUMENTS"), 0o644); err != nil { + t.Fatal(err) + } + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{Disable: true, CLI: []string{cliFile}}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + if _, ok := reg.Lookup("review"); ok { + t.Error("/review should NOT be registered under --no-prompt-templates") + } + if _, ok := reg.Lookup("cli"); ok { + t.Error("/cli should NOT be registered under --no-prompt-templates") + } + // Built-in slash commands are unaffected: the registry is non-empty. + if len(reg.List()) == 0 { + t.Error("built-in slash commands should still be registered under --no-prompt-templates") + } +} + +func TestBuildSlashRegistryGlobalOverridesCLI(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + // global prompt (TierGlobal) under ~/.pigo/prompts. + testutil.WritePrompt(t, home, "prompts", "dup.md", "FROM GLOBAL") + // CLI-tier file of the same name at the home root (not under prompts/). + cliFile := filepath.Join(home, "dup.md") + if err := os.WriteFile(cliFile, []byte("FROM CLI"), 0o644); err != nil { + t.Fatal(err) + } + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{CLI: []string{cliFile}}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + cmd, ok := reg.Lookup("dup") + if !ok { + t.Fatal("/dup not found") + } + if got := cmd.Expand(""); got != "FROM GLOBAL" { + t.Errorf("global should override CLI, got %q", got) + } + found := false + for _, e := range reg.Shadowed() { + if e.Name == "dup" && e.Tier == runtime.TierCLI { + found = true + } + } + if !found { + t.Errorf("CLI dup should be shadowed with TierCLI, got %v", reg.Shadowed()) + } +} diff --git a/pigo/internal/cli/prompts/prompts_config_test.go b/pigo/internal/cli/prompts/prompts_config_test.go new file mode 100644 index 0000000..ad45af2 --- /dev/null +++ b/pigo/internal/cli/prompts/prompts_config_test.go @@ -0,0 +1,115 @@ +package prompts + +// Tests for settings-tier prompt templates (US-007, #338): LoadPromptPaths +// loads file/dir entries (warning on missing), and BuildSlashRegistry +// registers them at the settings tier (overridden by global same-name +// templates). The applyFileConfig parse test stays in package main +// (cmd/pigo/prompts_config_test.go) since it drives cliOptions. + +import ( + "os" + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/testutil" + "github.com/smallnest/pigo/internal/runtime" +) + +func TestLoadSettingsPromptsFileDirMissing(t *testing.T) { + home := t.TempDir() + // file entry -> 1 cmd. + filePath := filepath.Join(home, "single.md") + if err := os.WriteFile(filePath, []byte("Single: $ARGUMENTS"), 0o644); err != nil { + t.Fatal(err) + } + // dir entry -> 2 cmds. + dirPath := filepath.Join(home, "promptsdir") + if err := os.MkdirAll(dirPath, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dirPath, "a.md"), []byte("A: $1"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dirPath, "b.md"), []byte("B: $1"), 0o644); err != nil { + t.Fatal(err) + } + // missing entry -> 0 cmds (warned, not fatal). + missing := filepath.Join(home, "nope") + + cmds := LoadPromptPaths([]string{filePath, dirPath, missing}) + if len(cmds) != 3 { + t.Fatalf("got %d cmds, want 3 (file=1 + dir=2 + missing=0)", len(cmds)) + } + names := map[string]bool{} + for _, c := range cmds { + names[c.Name] = true + } + for _, want := range []string{"single", "a", "b"} { + if !names[want] { + t.Errorf("missing cmd %q; got %v", want, names) + } + } +} + +func TestBuildSlashRegistryLoadsSettingsPrompts(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + // A settings-tier prompt dir (loaded via configPrompts). + settingsDir := filepath.Join(home, "settings-prompts") + if err := os.MkdirAll(settingsDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(settingsDir, "review.md"), []byte("FROM SETTINGS"), 0o644); err != nil { + t.Fatal(err) + } + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{Settings: []string{settingsDir}}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + cmd, ok := reg.Lookup("review") + if !ok { + t.Fatal("/review not found") + } + if got := cmd.Expand(""); got != "FROM SETTINGS" { + t.Errorf("settings prompt: got %q, want \"FROM SETTINGS\"", got) + } +} + +func TestBuildSlashRegistryGlobalOverridesSettings(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + // global prompt (TierGlobal) under ~/.pigo/prompts. + testutil.WritePrompt(t, home, "prompts", "dup.md", "FROM GLOBAL") + // settings-tier prompt (file) of the same name, placed at the home root + // (not under prompts/, so the global loop does not also load it). + settingsFile := filepath.Join(home, "dup.md") + if err := os.WriteFile(settingsFile, []byte("FROM SETTINGS"), 0o644); err != nil { + t.Fatal(err) + } + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{Settings: []string{settingsFile}}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + cmd, ok := reg.Lookup("dup") + if !ok { + t.Fatal("/dup not found") + } + if got := cmd.Expand(""); got != "FROM GLOBAL" { + t.Errorf("global should override settings, got %q", got) + } + // The settings loser is shadowed with TierSettings. + found := false + for _, e := range reg.Shadowed() { + if e.Name == "dup" && e.Tier == runtime.TierSettings { + found = true + } + } + if !found { + t.Errorf("settings dup should be shadowed with TierSettings, got %v", reg.Shadowed()) + } +} diff --git a/pigo/internal/cli/prompts/prompts_dir_test.go b/pigo/internal/cli/prompts/prompts_dir_test.go new file mode 100644 index 0000000..c481257 --- /dev/null +++ b/pigo/internal/cli/prompts/prompts_dir_test.go @@ -0,0 +1,91 @@ +package prompts + +// Tests for global prompt-template discovery (US-005, #336): BuildSlashRegistry +// loads both the legacy ~/.pigo/commands and the pi-aligned ~/.pigo/prompts +// (non-recursive, global tier), and a same-named template in prompts/ overrides +// the one in commands/ (last-write-wins within the global tier). + +import ( + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/testutil" +) + +// TestBuildSlashRegistryLoadsLegacyCommandsDir verifies the legacy +// ~/.pigo/commands directory still loads templates (regression). +func TestBuildSlashRegistryLoadsLegacyCommandsDir(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + testutil.WritePrompt(t, home, "commands", "legacy.md", "Legacy: $ARGUMENTS") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, PromptTemplateSources{}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + out, err := reg.ResolveOutcome("/legacy hi") + if err != nil { + t.Fatalf("ResolveOutcome: %v", err) + } + if !out.Handled || out.Prompt != "Legacy: hi" { + t.Errorf("/legacy = handled=%v prompt=%q, want handled=true \"Legacy: hi\"", out.Handled, out.Prompt) + } +} + +// TestBuildSlashRegistryLoadsPromptsDir verifies the pi-aligned ~/.pigo/prompts +// directory loads templates. +func TestBuildSlashRegistryLoadsPromptsDir(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + testutil.WritePrompt(t, home, "prompts", "review.md", "Review: $ARGUMENTS") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, PromptTemplateSources{}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + out, err := reg.ResolveOutcome("/review diff") + if err != nil { + t.Fatalf("ResolveOutcome: %v", err) + } + if !out.Handled || out.Prompt != "Review: diff" { + t.Errorf("/review = handled=%v prompt=%q, want handled=true \"Review: diff\"", out.Handled, out.Prompt) + } +} + +// TestBuildSlashRegistryPromptsOverridesCommands verifies that a same-named +// template in prompts/ overrides one in commands/ (both global tier; prompts is +// loaded second so last-write-wins), with no shadow entry. +func TestBuildSlashRegistryPromptsOverridesCommands(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + testutil.WritePrompt(t, home, "commands", "dup.md", "FROM COMMANDS") + testutil.WritePrompt(t, home, "prompts", "dup.md", "FROM PROMPTS") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, PromptTemplateSources{}) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + cmd, ok := reg.Lookup("dup") + if !ok { + t.Fatal("/dup not found") + } + if got := cmd.Expand(""); got != "FROM PROMPTS" { + t.Errorf("prompts should override commands on same name, got %q", got) + } + if len(reg.Shadowed()) != 0 { + t.Errorf("same-tier override must not shadow, got %v", reg.Shadowed()) + } +} + +// TestBuildSlashRegistryMissingDirsNoError verifies that with neither commands/ +// nor prompts/ present, BuildSlashRegistry returns no error (built-ins only). +func TestBuildSlashRegistryMissingDirsNoError(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, PromptTemplateSources{}) + if err != nil { + t.Fatalf("BuildSlashRegistry with no prompt dirs: %v", err) + } + if reg == nil { + t.Fatal("registry is nil") + } +} diff --git a/pigo/internal/cli/prompts/prompts_project_test.go b/pigo/internal/cli/prompts/prompts_project_test.go new file mode 100644 index 0000000..b53da09 --- /dev/null +++ b/pigo/internal/cli/prompts/prompts_project_test.go @@ -0,0 +1,139 @@ +package prompts + +// Tests for project-level .pigo/prompts (US-006, #337): loaded at the project +// tier only when the project is trusted, overrides global same-name templates, +// and is suppressed by --no-prompt-templates. + +import ( + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/testutil" + "github.com/smallnest/pigo/internal/runtime" +) + +// TestBuildSlashRegistryLoadsProjectPromptsTrusted: with the project trusted, +// .pigo/prompts/*.md loads at the project tier. +func TestBuildSlashRegistryLoadsProjectPromptsTrusted(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) // empty global + cwdTmp := t.TempDir() + testutil.WritePrompt(t, cwdTmp, filepath.Join(".pigo", "prompts"), "review.md", "Review: $ARGUMENTS") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{ + ProjectDir: filepath.Join(cwdTmp, ".pigo", "prompts"), + ProjectTrusted: true, + }) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + out, err := reg.ResolveOutcome("/review diff") + if err != nil { + t.Fatalf("ResolveOutcome: %v", err) + } + if !out.Handled || out.Prompt != "Review: diff" { + t.Errorf("/review = handled=%v prompt=%q, want \"Review: diff\"", out.Handled, out.Prompt) + } +} + +// TestBuildSlashRegistryProjectPromptsUntrustedSkipped: when the project is not +// trusted, .pigo/prompts is not loaded. +func TestBuildSlashRegistryProjectPromptsUntrustedSkipped(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + cwdTmp := t.TempDir() + testutil.WritePrompt(t, cwdTmp, filepath.Join(".pigo", "prompts"), "review.md", "Review: $ARGUMENTS") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{ + ProjectDir: filepath.Join(cwdTmp, ".pigo", "prompts"), + ProjectTrusted: false, + }) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + if _, ok := reg.Lookup("review"); ok { + t.Error("/review should NOT load from an untrusted project") + } +} + +// TestBuildSlashRegistryProjectMissingDirNoError: a missing .pigo/prompts is +// not an error (most projects don't have one). +func TestBuildSlashRegistryProjectMissingDirNoError(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + cwdTmp := t.TempDir() // no .pigo/prompts created + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{ + ProjectDir: filepath.Join(cwdTmp, ".pigo", "prompts"), + ProjectTrusted: true, + }) + if err != nil { + t.Fatalf("missing .pigo/prompts should not error, got %v", err) + } + if reg == nil { + t.Fatal("registry is nil") + } +} + +// TestBuildSlashRegistryProjectOverridesGlobal: a project template overrides a +// same-named global one (project tier wins, global shadowed). +func TestBuildSlashRegistryProjectOverridesGlobal(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + // global + testutil.WritePrompt(t, home, "prompts", "dup.md", "FROM GLOBAL") + // project + cwdTmp := t.TempDir() + testutil.WritePrompt(t, cwdTmp, filepath.Join(".pigo", "prompts"), "dup.md", "FROM PROJECT") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{ + ProjectDir: filepath.Join(cwdTmp, ".pigo", "prompts"), + ProjectTrusted: true, + }) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + cmd, ok := reg.Lookup("dup") + if !ok { + t.Fatal("/dup not found") + } + if got := cmd.Expand(""); got != "FROM PROJECT" { + t.Errorf("project should override global, got %q", got) + } + found := false + for _, e := range reg.Shadowed() { + if e.Name == "dup" && e.Tier == runtime.TierGlobal { + found = true + } + } + if !found { + t.Errorf("global dup should be shadowed with TierGlobal, got %v", reg.Shadowed()) + } +} + +// TestBuildSlashRegistryNoPromptTemplatesDisablesProject: --no-prompt-templates +// suppresses project prompts too. +func TestBuildSlashRegistryNoPromptTemplatesDisablesProject(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + cwdTmp := t.TempDir() + testutil.WritePrompt(t, cwdTmp, filepath.Join(".pigo", "prompts"), "review.md", "Review: $ARGUMENTS") + + reg, err := BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, + PromptTemplateSources{ + Disable: true, + ProjectDir: filepath.Join(cwdTmp, ".pigo", "prompts"), + ProjectTrusted: true, + }) + if err != nil { + t.Fatalf("BuildSlashRegistry: %v", err) + } + if _, ok := reg.Lookup("review"); ok { + t.Error("/review should NOT load under --no-prompt-templates") + } +} diff --git a/pigo/internal/cli/prompts/registry.go b/pigo/internal/cli/prompts/registry.go new file mode 100644 index 0000000..fcfa077 --- /dev/null +++ b/pigo/internal/cli/prompts/registry.go @@ -0,0 +1,387 @@ +// Package prompts holds the slash-command registry assembly shared by the REPL +// (internal/cli/repl) and the forthcoming TUI (internal/cli/tui). It was sunk +// out of the repl package (#383) so both front-ends wire the same built-in, +// live-state, plugin-declared, prompt-template and skill commands from one +// owner, avoiding drift between the two command surfaces. +// +// The logic here is a verbatim move of repl's former private +// buildSlashRegistry/loadPromptPaths/promptTemplateSources (plus their +// register helpers), exported unchanged so REPL behavior is identical. +package prompts + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" +) + +// PromptTemplateSources carries the prompt-template discovery sources that +// BuildSlashRegistry loads beyond the global ~/.pigo/{commands,prompts} dirs. +// Settings is the config.toml `prompts` array (TierSettings); CLI is the +// --prompt-template flag list (TierCLI, wired in #339). Each entry is a file or +// directory (loaded non-recursively). Missing paths are warned and skipped. +type PromptTemplateSources struct { + Settings []string + CLI []string + // Disable (--no-prompt-templates) turns off all prompt-template discovery + // (global, project, settings, CLI); built-ins and skills are unaffected. + Disable bool + // ProjectDir is the project-local prompts dir (.pigo/prompts in the working + // dir), loaded at the project tier only when ProjectTrusted is true. + ProjectDir string + // ProjectTrusted reports whether the working directory is trusted; project + // templates load only then (mirrors pi: project prompts after the project is + // trusted). + ProjectTrusted bool +} + +// LoadPromptPaths loads prompt templates from each path (file or dir), skipping +// and warning on paths that don't exist or fail to read. It is tier-agnostic; +// the caller registers each result at the desired tier (AddSettings/AddCLI). +func LoadPromptPaths(paths []string) []runtime.SlashCommand { + var out []runtime.SlashCommand + for _, p := range paths { + info, err := os.Stat(p) + if err != nil { + fmt.Fprintf(os.Stderr, "pigo: prompts path %q not found, skipping\n", p) + continue + } + var cmds []runtime.SlashCommand + if info.IsDir() { + cmds, err = runtime.LoadUserCommandsDir(p) + } else { + c, e := runtime.LoadPromptFile(p) + if e != nil { + err = e + } else { + cmds = []runtime.SlashCommand{c} + } + } + if err != nil { + fmt.Fprintf(os.Stderr, "pigo: prompts path %q: %v\n", p, err) + continue + } + out = append(out, cmds...) + } + return out +} + +// BuildSlashRegistry assembles the slash-command registry: compile-time +// built-ins seeded by runtime.NewSlashRegistry, the live-state action commands +// (/model, /help) bound to live, user declarative templates loaded from +// ~/.pigo/commands (or $PIGO_HOME/commands), plugin-declared commands from the +// loaded Manager, plus the pre-loaded skills — each surfaced as a "/skill-name" +// command (mirrors Claude Code's /skill invocation). A missing directory is not an +// error. Names that collide with a built-in are shadowed (the built-in wins) and +// reported on stderr. The skills slice is loaded once by setupAgentEnv (empty +// under --no-skills), so no /skill-name commands are registered when it is +// empty. mgr may be nil (no plugins loaded). +func BuildSlashRegistry(live *cli.LiveConfig, skills []*runtime.Skill, mgr *plugin.Manager, srcs PromptTemplateSources) (*runtime.SlashRegistry, error) { + reg := runtime.NewSlashRegistry() + RegisterLiveCommands(reg, live) + RegisterPluginCommands(reg, mgr) + // --no-prompt-templates disables all prompt-template discovery (global, + // settings, CLI); built-in slash commands and skills are unaffected. + if !srcs.Disable { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return reg, nil // built-ins only + } + dir = filepath.Join(home, ".pigo") + } + // Load user prompt templates from both the legacy ~/.pigo/commands and + // the pi-aligned ~/.pigo/prompts (both non-recursive, global tier). + // Loading commands first means a same-named template in prompts/ + // overrides the legacy one (last-write-wins within the global tier). A + // missing directory is not an error (LoadUserCommandsDir returns nil, + // nil for IsNotExist). + for _, sub := range []string{"commands", "prompts"} { + cmds, err := runtime.LoadUserCommandsDir(filepath.Join(dir, sub)) + if err != nil { + return reg, err + } + for _, c := range cmds { + reg.AddUser(c) + } + } + // Settings-tier templates from the config.toml `prompts` array, then + // CLI-tier templates from --prompt-template. Each entry is a file or + // dir; missing paths are warned and skipped. + for _, c := range LoadPromptPaths(srcs.Settings) { + reg.AddSettings(c) + } + for _, c := range LoadPromptPaths(srcs.CLI) { + reg.AddCLI(c) + } + // Project-tier templates from .pigo/prompts in the working directory, + // loaded only when the project is trusted (mirrors pi). A missing dir is + // not an error. Overrides global/settings/CLI (project tier is higher). + if srcs.ProjectTrusted && srcs.ProjectDir != "" { + cmds, err := runtime.LoadUserCommandsDir(srcs.ProjectDir) + if err != nil { + return reg, err + } + for _, c := range cmds { + reg.AddProject(c) + } + } + } + // Register skills as /skill-name commands from the pre-loaded set (shared with + // prompt injection in setupAgentEnv, so the directory is read once). All + // skills — including disable-model-invocation ones — get a slash command; the + // prompt-injection side filters the disabled ones. Under --no-skills the set + // is empty, so nothing is registered. + for _, s := range skills { + reg.AddSkill(s.SlashCommand()) + } + if sh := reg.Shadowed(); len(sh) > 0 { + parts := make([]string, len(sh)) + for i, e := range sh { + parts[i] = e.String() + } + fmt.Fprintf(os.Stderr, "pigo: commands shadowed by higher-priority source (rename to use): %v\n", parts) + } + return reg, nil +} + +// RegisterPluginCommands installs each plugin-declared slash command +// (Manager.Commands()) into the registry as a hybrid (Run) command. Invoking it +// RPCs the owning plugin (Plugin.CallCommand), returns the plugin's +// notifications as the outcome Message, and returns the plugin's Prompt to run +// as the next turn. Plugin commands are registered with AddPlugin so a same-named +// built-in still wins (existing precedence preserved) and a collision is +// reported as shadowed. mgr may be nil (no plugins), in which case this is a +// no-op. +// +// The args passed to CallCommand are the invocation's raw argument text encoded +// as a JSON string (json.RawMessage of a quoted string), never null: the host +// (node #263) expects a JSON string for a no-arg command, so a bare "/cmd" +// sends `""` rather than nil. Each command captures its own plugin and spec name +// (loop variables copied per-iteration). +func RegisterPluginCommands(reg *runtime.SlashRegistry, mgr *plugin.Manager) { + if mgr == nil { + return + } + for _, pc := range mgr.Commands() { + pc := pc // capture per iteration + reg.AddPlugin(runtime.SlashCommand{ + Name: pc.Spec.Name, + Description: pc.Spec.Description, + Run: func(args string) (message, prompt string) { + // Encode the raw arg text as a JSON string ("" for no args), matching + // the host's CommandCallParams.Args contract (a JSON string, never + // null). json.Marshal of a Go string always succeeds. + raw, _ := json.Marshal(args) + res, err := pc.Plugin.CallCommand(context.Background(), pc.Spec.Name, json.RawMessage(raw)) + if err != nil { + return fmt.Sprintf("plugin command %q failed: %v", pc.Spec.Name, err), "" + } + return formatNotifications(res.Notifications), res.Prompt + }, + }) + } +} + +// formatNotifications renders a plugin command's notifications into a single +// block to surface to the user, one per line, prefixed by their type (when set) +// so severity is visible. Returns "" when there are none. +func formatNotifications(notes []plugin.CommandNotification) string { + if len(notes) == 0 { + return "" + } + var b strings.Builder + for i, n := range notes { + if i > 0 { + b.WriteString("\n") + } + if n.Type != "" { + b.WriteString("[") + b.WriteString(n.Type) + b.WriteString("] ") + } + b.WriteString(n.Message) + } + return b.String() +} + +// RegisterLiveCommands installs the built-in action commands that need live +// runtime state. /model views or switches the active model; /help lists the +// available commands. These are instance built-ins (AddBuiltin) because their +// closures must capture live and the registry — state unreachable from an +// init()-time global registration. +func RegisterLiveCommands(reg *runtime.SlashRegistry, live *cli.LiveConfig) { + reg.AddBuiltin(runtime.SlashCommand{ + Name: "model", + Description: "view or switch the active model: /model [model-id] (see /models for presets)", + Action: func(args string) string { + id := strings.TrimSpace(args) + if id == "" { + return fmt.Sprintf("model: %s (provider: %s)\nrun /models to see presets, or /model to switch", live.Model, live.ProviderName) + } + prov, providerName, err := provider.ResolveProvider(id, live.BaseURL, live.Protocol, "", os.Getenv) + if err != nil { + return fmt.Sprintf("model: cannot switch to %q: %v", id, err) + } + live.Model = id + live.ProviderName = providerName + live.Provider = prov + return fmt.Sprintf("model switched to %s (provider: %s)", id, providerName) + }, + }) + reg.AddBuiltin(runtime.SlashCommand{ + Name: "models", + Description: "list preset providers and models you can switch to", + Action: func(args string) string { return presetListing(strings.TrimSpace(args)) }, + }) + // thinkAction views or switches the reasoning-effort level. It backs both + // /think and its alias /effect, so the two commands share identical behavior. + thinkAction := func(args string) string { + lvl := strings.TrimSpace(args) + if lvl == "" { + cur := live.ThinkingLevel + if cur == "" { + cur = agentcore.ThinkingOff + } + return fmt.Sprintf("think: %s\nswitch with /think ", cur) + } + v, ok := validThinkingLevel(lvl) + if !ok { + return fmt.Sprintf("think: invalid level %q (want off|minimal|low|medium|high|xhigh|max)", lvl) + } + live.ThinkingLevel = v + return fmt.Sprintf("think level set to %s (applies to the next turn)", v) + } + reg.AddBuiltin(runtime.SlashCommand{ + Name: "think", + ArgumentHint: "[off|minimal|low|medium|high|xhigh|max]", + Description: "view or switch the reasoning-effort level; takes effect on the next turn", + Action: thinkAction, + }) + reg.AddBuiltin(runtime.SlashCommand{ + Name: "effect", + ArgumentHint: "[off|minimal|low|medium|high|xhigh|max]", + Description: "alias of /think: view or switch the reasoning-effort level", + Action: thinkAction, + }) + reg.AddBuiltin(runtime.SlashCommand{ + Name: "help", + Description: "list available slash commands", + Action: func(string) string { + color := ui.Enabled() + var b strings.Builder + b.WriteString(ui.Colorize(color, ui.Bold, "available commands:")) + for _, c := range reg.List() { + b.WriteString("\n ") + b.WriteString(ui.Colorize(color, ui.Cyan, "/"+c.Name)) + rest := "" + if c.ArgumentHint != "" { + rest += " " + c.ArgumentHint + } + if c.Description != "" { + rest += " - " + c.Description + } + rest += " (source: " + c.Tier.String() + ")" + b.WriteString(ui.Colorize(color, ui.Dim, rest)) + } + return b.String() + }, + }) + // /exit, /quit, /compact, /fork, /clone, /tree, /export, /import, /copy, + // /session and /status are intercepted by the REPL loop before slash resolution + // (they must return from the loop, run an agent stream, or read/swap the active + // session/leaf — none of which an Action closure can do). They are registered + // here only so /help lists them; their Action is never actually reached. + for _, c := range []struct{ name, desc string }{ + {"exit", "exit the REPL"}, + {"quit", "exit the REPL"}, + {"compact", "summarize and compact the conversation context now"}, + {"fork", "branch from a historical message into a new session: /fork [n]"}, + {"clone", "duplicate the current session into an independent branch"}, + {"tree", "show the session branch tree; switch active branch: /tree [n]"}, + {"rewind", "roll files and the conversation back to before an earlier turn: /rewind [n]"}, + {"export", "export the session to a file: /export [path.jsonl|path.html]"}, + {"import", "import a JSONL export as a new session: /import "}, + {"copy", "copy the most recent assistant reply to the clipboard"}, + {"session", "show session stats: messages, tokens, model, compactions"}, + {"status", "show session status: runtime config, context, telemetry, credentials, environment"}, + {"goal", "run autonomously toward a goal: /goal [--tokens N] | pause | resume | clear"}, + {"btw", "ask a quick side question without touching the main conversation: /btw (bare /btw reopens the last one)"}, + {"dream", "consolidate memory now (dedupe, merge, prune, distill); /dream --dry-run previews without writing"}, + {"remote-control", "mirror this session to a phone/browser on your LAN: /remote-control [stop|status]"}, + } { + reg.AddBuiltin(runtime.SlashCommand{ + Name: c.name, + Description: c.desc, + Action: func(string) string { return "" }, + }) + } +} + +// validThinkingLevel reports whether s is one of the known reasoning-effort +// levels and returns the typed value. It mirrors the enum in agentcore so a +// /think argument can be validated without importing the config layer. +func validThinkingLevel(s string) (agentcore.ThinkingLevel, bool) { + switch agentcore.ThinkingLevel(s) { + case agentcore.ThinkingOff, agentcore.ThinkingMinimal, agentcore.ThinkingLow, + agentcore.ThinkingMedium, agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax: + return agentcore.ThinkingLevel(s), true + default: + return "", false + } +} + +// presetListing renders the preset provider/model catalog for /models. With an +// argument it filters to a single provider (e.g. "/models nvidia"). Providers +// are grouped and shown with the env var their API key is read from (referenced +// by name only, never a value). The output guides the user to `/model `. +func presetListing(filter string) string { + var b strings.Builder + b.WriteString("preset providers & models (switch with /model ):") + shown := 0 + for _, pv := range provider.PresetProviders { + if filter != "" && !strings.EqualFold(filter, pv.Name) { + continue + } + models := provider.PresetsByProvider(pv.Name) + if len(models) == 0 { + continue + } + shown++ + b.WriteString("\n\n") + b.WriteString(pv.Name) + if pv.EnvVar != "" { + b.WriteString(" (API key: $") + b.WriteString(pv.EnvVar) + b.WriteString(")") + } else { + b.WriteString(" (local, no API key)") + } + for _, m := range models { + b.WriteString("\n ") + b.WriteString(m.ID) + if m.DisplayName != "" { + b.WriteString(" — ") + b.WriteString(m.DisplayName) + } + } + } + if shown == 0 { + if filter != "" { + return fmt.Sprintf("no preset provider named %q (try openrouter, nvidia, or ollama)", filter) + } + return "no presets configured" + } + return b.String() +} diff --git a/pigo/internal/cli/prompts/think_test.go b/pigo/internal/cli/prompts/think_test.go new file mode 100644 index 0000000..a8b19f9 --- /dev/null +++ b/pigo/internal/cli/prompts/think_test.go @@ -0,0 +1,87 @@ +package prompts + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/runtime" +) + +// TestThinkCommandSwitchesLevel verifies /think mutates the live thinking level +// so the next turn picks it up, and that a bare /think reports the current level. +func TestThinkCommandSwitchesLevel(t *testing.T) { + live := &cli.LiveConfig{Model: "test", ProviderName: "test", ThinkingLevel: agentcore.ThinkingMedium} + reg := runtime.NewSlashRegistry() + RegisterLiveCommands(reg, live) + + out, err := reg.ResolveOutcome("/think high") + if err != nil { + t.Fatalf("ResolveOutcome /think high: %v", err) + } + if live.ThinkingLevel != agentcore.ThinkingHigh { + t.Errorf("ThinkingLevel = %q, want high", live.ThinkingLevel) + } + if !strings.Contains(out.Message, "high") { + t.Errorf("message = %q, want it to mention high", out.Message) + } + + // Bare /think reports the current level without changing it. + out, err = reg.ResolveOutcome("/think") + if err != nil { + t.Fatalf("ResolveOutcome /think: %v", err) + } + if live.ThinkingLevel != agentcore.ThinkingHigh { + t.Errorf("bare /think mutated level to %q", live.ThinkingLevel) + } + if !strings.Contains(out.Message, "high") { + t.Errorf("bare /think message = %q, want current level high", out.Message) + } +} + +// TestThinkCommandRejectsInvalid verifies an unknown level is rejected and the +// live level is left unchanged. +func TestThinkCommandRejectsInvalid(t *testing.T) { + live := &cli.LiveConfig{Model: "test", ProviderName: "test", ThinkingLevel: agentcore.ThinkingLow} + reg := runtime.NewSlashRegistry() + RegisterLiveCommands(reg, live) + + out, err := reg.ResolveOutcome("/think bogus") + if err != nil { + t.Fatalf("ResolveOutcome /think bogus: %v", err) + } + if live.ThinkingLevel != agentcore.ThinkingLow { + t.Errorf("invalid level changed ThinkingLevel to %q", live.ThinkingLevel) + } + if !strings.Contains(out.Message, "invalid") { + t.Errorf("message = %q, want an invalid-level notice", out.Message) + } +} + +// TestEffectAliasesThink verifies /effect behaves identically to /think: it +// switches the live thinking level and a bare /effect reports the current level. +func TestEffectAliasesThink(t *testing.T) { + live := &cli.LiveConfig{Model: "test", ProviderName: "test", ThinkingLevel: agentcore.ThinkingMedium} + reg := runtime.NewSlashRegistry() + RegisterLiveCommands(reg, live) + + out, err := reg.ResolveOutcome("/effect high") + if err != nil { + t.Fatalf("ResolveOutcome /effect high: %v", err) + } + if live.ThinkingLevel != agentcore.ThinkingHigh { + t.Errorf("ThinkingLevel = %q, want high", live.ThinkingLevel) + } + if !strings.Contains(out.Message, "high") { + t.Errorf("message = %q, want it to mention high", out.Message) + } + + out, err = reg.ResolveOutcome("/effect") + if err != nil { + t.Fatalf("ResolveOutcome /effect: %v", err) + } + if !strings.Contains(out.Message, "high") { + t.Errorf("bare /effect message = %q, want current level high", out.Message) + } +} diff --git a/pigo/internal/cli/providerhelp.go b/pigo/internal/cli/providerhelp.go new file mode 100644 index 0000000..b2697ee --- /dev/null +++ b/pigo/internal/cli/providerhelp.go @@ -0,0 +1,31 @@ +// This file holds ProviderHelp, the "--help" provider block moved from cmd/pigo +// (US-004, #361). It enumerates the built-in provider registry so the values +// accepted by --provider (and their env vars / default base URLs / protocols) +// stay in sync with the code rather than being hand-maintained. +package cli + +import ( + "fmt" + "io" + "strings" + + "github.com/smallnest/pigo/internal/provider" +) + +// PrintProviderHelp writes the "Supported providers" block appended to `--help` +// output. It enumerates the built-in provider registry so the list of values +// accepted by --provider (and their env vars / default base URLs / protocols) +// stays in sync with the code rather than being hand-maintained. +func PrintProviderHelp(w io.Writer) { + fmt.Fprintf(w, "\nSupported --provider names (name: ENV_VARS -> default base URL [protocol]):\n") + for _, spec := range provider.ProviderSpecs() { + base := spec.DefaultBaseURL + if strings.TrimSpace(base) == "" { + base = "(composed from env)" + } + fmt.Fprintf(w, " %s: %s -> %s [%s]\n", + spec.Name, strings.Join(spec.EnvVars, ", "), base, spec.Protocol) + } + fmt.Fprintf(w, "\nBase URL override precedence: --base-url > -specific *_BASE_URL env > generic _BASE_URL env > registry default.\n") + fmt.Fprintf(w, "API key env fallback: any provider also accepts the generic _API_KEY convention.\n") +} diff --git a/pigo/internal/cli/providerhelp_test.go b/pigo/internal/cli/providerhelp_test.go new file mode 100644 index 0000000..d950d3c --- /dev/null +++ b/pigo/internal/cli/providerhelp_test.go @@ -0,0 +1,32 @@ +package cli + +import ( + "bytes" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/provider" +) + +// TestPrintProviderHelp verifies the --help "Supported providers" block is +// derived from the registry: every built-in provider name, its env vars, and +// its protocol must appear, so the docs cannot silently drift from the code. +func TestPrintProviderHelp(t *testing.T) { + var buf bytes.Buffer + PrintProviderHelp(&buf) + out := buf.String() + + for _, spec := range provider.ProviderSpecs() { + if !strings.Contains(out, spec.Name+":") { + t.Errorf("help output missing provider %q", spec.Name) + } + for _, env := range spec.EnvVars { + if !strings.Contains(out, env) { + t.Errorf("help output missing env var %q for provider %q", env, spec.Name) + } + } + if !strings.Contains(out, "["+spec.Protocol+"]") { + t.Errorf("help output missing protocol %q for provider %q", spec.Protocol, spec.Name) + } + } +} diff --git a/pigo/internal/cli/repl/autocomplete_label_test.go b/pigo/internal/cli/repl/autocomplete_label_test.go new file mode 100644 index 0000000..c552cfb --- /dev/null +++ b/pigo/internal/cli/repl/autocomplete_label_test.go @@ -0,0 +1,62 @@ +package repl + +// Tests for formatSlashAutocompleteLabel (US-010, #340): the Tab-completion +// label renders as "name - description", omitting the hint +// segment when absent and falling back to the first body line for description. + +import ( + "bufio" + "io" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/runtime" +) + +func TestFormatSlashAutocompleteLabel(t *testing.T) { + cases := []struct { + name string + cmd runtime.SlashCommand + want string + }{ + {"hint+desc", runtime.SlashCommand{Name: "review", ArgumentHint: "", Description: "Review PRs"}, "review - Review PRs"}, + {"desc only", runtime.SlashCommand{Name: "review", Description: "Review PRs"}, "review - Review PRs"}, + {"hint only", runtime.SlashCommand{Name: "wr", ArgumentHint: "[instructions]"}, "wr [instructions]"}, + {"neither", runtime.SlashCommand{Name: "model"}, "model"}, + } + for _, c := range cases { + if got := formatSlashAutocompleteLabel(c.cmd); got != c.want { + t.Errorf("%s: got %q, want %q", c.name, got, c.want) + } + } +} + +// TestFormatSlashAutocompleteLabelDescriptionFallback verifies the description +// fallback from #334 (first non-empty body line) flows through to the label. +func TestFormatSlashAutocompleteLabelDescriptionFallback(t *testing.T) { + cmd, err := runtime.ParseUserCommand("bare", []byte("First line is the desc\nbody")) + if err != nil { + t.Fatal(err) + } + if got := formatSlashAutocompleteLabel(cmd); got != "bare - First line is the desc" { + t.Errorf("fallback label = %q, want \"bare - First line is the desc\"", got) + } +} + +// TestSlashAutocompleteSuggestionStillCompletesName verifies that a slash +// command with an argument-hint is still completable by name (Tab inserts +// "/name"); the label is a display annotation, not the inserted text. +func TestSlashAutocompleteSuggestionStillCompletesName(t *testing.T) { + reg := runtime.NewSlashRegistry() + reg.AddBuiltin(runtime.SlashCommand{Name: "model", Action: func(string) string { return "" }}) + reg.AddUser(runtime.SlashCommand{ + Name: "review", + ArgumentHint: "", + Description: "Review PRs", + Expand: func(string) string { return "" }, + }) + e := newREPLLineEditor(strings.NewReader(""), bufio.NewReader(strings.NewReader("")), io.Discard, reg, nil) + if got := e.suggestion("/rev"); got != "/review" { + t.Errorf("suggestion(/rev) = %q, want /review (name, not the label)", got) + } +} diff --git a/pigo/internal/cli/repl/btw_help_test.go b/pigo/internal/cli/repl/btw_help_test.go new file mode 100644 index 0000000..be305b6 --- /dev/null +++ b/pigo/internal/cli/repl/btw_help_test.go @@ -0,0 +1,58 @@ +package repl + +// Tests for /btw discoverability (#283, US-006/FR-10): /btw must appear in the +// /help listing and be completable at the REPL slash prompt. Both flow from +// registering "btw" as a listed built-in in registerLiveCommands; these tests +// pin that so the command can't silently drop off help/completion. + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/runtime" +) + +// TestBtwListedInHelp verifies registerLiveCommands registers /btw so /help (and +// any consumer of reg.List()) surfaces it with a usage description. +func TestBtwListedInHelp(t *testing.T) { + reg := runtime.NewSlashRegistry() + prompts.RegisterLiveCommands(reg, &cli.LiveConfig{}) + + var btw *runtime.SlashCommand + for _, c := range reg.List() { + if c.Name == "btw" { + cc := c + btw = &cc + break + } + } + if btw == nil { + t.Fatalf("registerLiveCommands must register /btw so /help lists it") + } + if !strings.Contains(btw.Description, "side question") { + t.Errorf("/btw help description should explain the side question, got %q", btw.Description) + } +} + +// TestBtwSlashCompletion verifies the line editor completes "/b" to "/btw" once +// the command is registered — the same reg.List() drives both help and +// completion, so registration is all that's needed. +func TestBtwSlashCompletion(t *testing.T) { + reg := runtime.NewSlashRegistry() + prompts.RegisterLiveCommands(reg, &cli.LiveConfig{}) + e := newREPLLineEditor(nil, nil, nil, reg, nil) + + cands := e.suggestions("/bt") + found := false + for _, c := range cands { + if c == "/btw" { + found = true + break + } + } + if !found { + t.Errorf("expected /btw among completions for %q, got %v", "/bt", cands) + } +} diff --git a/pigo/internal/cli/repl/btw_isolation_test.go b/pigo/internal/cli/repl/btw_isolation_test.go new file mode 100644 index 0000000..25a1f5d --- /dev/null +++ b/pigo/internal/cli/repl/btw_isolation_test.go @@ -0,0 +1,148 @@ +package repl + +// Isolation / zero-pollution tests for /btw (#284, PRD Success Metrics). These +// lock the feature's most important correctness guarantee: a side thread's +// question and answer NEVER touch the main conversation and NEVER hit disk. The +// tests drive the whole runREPL loop with the fake replProvider and assert, in +// one place, every observable that a leak would perturb: the main context's +// message count AND content, the session store on disk, and the persistence +// bookkeeping (deps.persisted / deps.curLeaf / deps.header.UpdatedAt). + +import ( + "bytes" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// snapshotMessages returns the flattened text of every message so a test can +// assert the main context is byte-for-byte unchanged, not merely the same +// length (a leak could replace a message without changing the count). Content +// is a field on the concrete message types, not on the Message interface, so we +// type-switch on the two kinds /btw could ever append. +func snapshotMessages(msgs agentcore.MessageList) []string { + out := make([]string, len(msgs)) + for i, m := range msgs { + var text string + switch v := m.(type) { + case agentcore.UserMessage: + text = agentcore.ContentToText(v.Content) + case agentcore.AssistantMessage: + text = agentcore.ContentToText(v.Content) + } + out[i] = m.Role() + ":" + text + } + return out +} + +// seedMainContext appends a real user+assistant exchange to the main context so +// the isolation tests start from a non-empty conversation — proving /btw leaves +// existing history untouched, not just that it avoids growing an empty slice. +func seedMainContext(deps *replDeps) { + deps.agentCtx.Messages = append(deps.agentCtx.Messages, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("main question")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("main answer")}}, + ) +} + +// TestBtwMainContextZeroGrowth verifies one /btw Q&A leaves the main context's +// message count AND content exactly as before (AC-1). +func TestBtwMainContextZeroGrowth(t *testing.T) { + p := &replProvider{reply: "side answer"} + deps, _ := newTestDeps(t, p) + seedMainContext(&deps) + before := snapshotMessages(deps.agentCtx.Messages) + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/btw a quick question?\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 1 { + t.Fatalf("expected exactly 1 side run, got %d", p.calls) + } + after := snapshotMessages(deps.agentCtx.Messages) + if len(after) != len(before) { + t.Fatalf("main context grew: %d → %d messages", len(before), len(after)) + } + for i := range before { + if after[i] != before[i] { + t.Fatalf("main context message %d changed:\n before %q\n after %q", i, before[i], after[i]) + } + } +} + +// TestBtwNoStoreWrites verifies /btw writes nothing to the session store: no +// entries are appended for the session on disk (AC-2). +func TestBtwNoStoreWrites(t *testing.T) { + p := &replProvider{reply: "answer"} + deps, store := newTestDeps(t, p) + seedMainContext(&deps) + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/btw does this persist?\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + // The main loop persists only on a real turn; /btw must trigger none. Since + // the seeded messages were never run through streamRun, the store should hold + // no entries for this session at all. + if _, entries, err := store.LoadEntries(deps.header.ID); err == nil && len(entries) != 0 { + t.Fatalf("/btw must not write to the store, got %d entries", len(entries)) + } +} + +// TestBtwFollowUpsLeaveContextUnchanged verifies that ≥3 follow-ups in the same +// side thread still leave the main context's count and content unchanged (AC-3). +func TestBtwFollowUpsLeaveContextUnchanged(t *testing.T) { + p := &replProvider{reply: "ok"} + deps, _ := newTestDeps(t, p) + seedMainContext(&deps) + before := snapshotMessages(deps.agentCtx.Messages) + + var out bytes.Buffer + // One /btw plus three bare follow-ups, then leave the thread and exit. + in := strings.NewReader("/btw first?\nsecond?\nthird?\nfourth?\n/exit\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 4 { + t.Fatalf("expected 4 side runs (1 + 3 follow-ups), got %d", p.calls) + } + after := snapshotMessages(deps.agentCtx.Messages) + if len(after) != len(before) { + t.Fatalf("main context grew across follow-ups: %d → %d", len(before), len(after)) + } + for i := range before { + if after[i] != before[i] { + t.Fatalf("follow-ups changed main message %d: %q → %q", i, before[i], after[i]) + } + } +} + +// TestBtwPersistenceBookkeepingUnchanged verifies /btw does not advance the +// persistence bookkeeping: deps.persisted, deps.curLeaf and deps.header.UpdatedAt +// are all identical before and after (AC-4). +func TestBtwPersistenceBookkeepingUnchanged(t *testing.T) { + p := &replProvider{reply: "x"} + deps, _ := newTestDeps(t, p) + seedMainContext(&deps) + // Give the bookkeeping non-zero starting values so the test would catch a + // reset-to-zero as well as an increment. + deps.persisted = 2 + deps.curLeaf = "leaf-abc" + beforeUpdated := deps.header.UpdatedAt + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/btw hi\nmore?\n/exit\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if deps.persisted != 2 { + t.Errorf("deps.persisted changed: 2 → %d", deps.persisted) + } + if deps.curLeaf != "leaf-abc" { + t.Errorf("deps.curLeaf changed: %q → %q", "leaf-abc", deps.curLeaf) + } + if !deps.header.UpdatedAt.Equal(beforeUpdated) { + t.Errorf("deps.header.UpdatedAt changed: %v → %v", beforeUpdated, deps.header.UpdatedAt) + } +} diff --git a/pigo/internal/cli/repl/btw_test.go b/pigo/internal/cli/repl/btw_test.go new file mode 100644 index 0000000..2d68948 --- /dev/null +++ b/pigo/internal/cli/repl/btw_test.go @@ -0,0 +1,171 @@ +package repl + +// Tests for the /btw side-thread command (#279): a side question must run an +// agent stream but MUST NOT mutate or persist the main conversation. These +// drive the whole runREPL loop with the fake replProvider, then assert the main +// context and persistence state are unchanged. + +import ( + "bytes" + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/btw" +) + +// TestBtwDoesNotPolluteMainContext verifies that "/btw " launches a run +// (provider called) yet appends nothing to deps.agentCtx.Messages. +func TestBtwDoesNotPolluteMainContext(t *testing.T) { + p := &replProvider{reply: "side answer"} + deps, _ := newTestDeps(t, p) + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/btw why pointers?\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 1 { + t.Fatalf("expected side question to launch exactly 1 run, got %d", p.calls) + } + if len(deps.agentCtx.Messages) != 0 { + t.Fatalf("main context must be untouched by /btw, got %d messages", len(deps.agentCtx.Messages)) + } + if !strings.Contains(out.String(), btw.BtwHeader) { + t.Errorf("expected side-thread header %q in output", btw.BtwHeader) + } + if !strings.Contains(out.String(), "side answer") { + t.Errorf("expected the side answer to be printed, got: %q", out.String()) + } +} + +// TestBtwDoesNotPersist verifies /btw writes nothing to disk: deps.persisted and +// deps.curLeaf are unchanged, and no session entries were appended. +func TestBtwDoesNotPersist(t *testing.T) { + p := &replProvider{reply: "answer"} + deps, store := newTestDeps(t, p) + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/btw quick q\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if deps.persisted != 0 { + t.Fatalf("deps.persisted must stay 0 after /btw, got %d", deps.persisted) + } + if deps.curLeaf != "" { + t.Fatalf("deps.curLeaf must stay empty after /btw, got %q", deps.curLeaf) + } + if _, entries, err := store.LoadEntries(deps.header.ID); err == nil && len(entries) != 0 { + t.Fatalf("no session entries should be persisted by /btw, got %d", len(entries)) + } +} + +// TestBtwBareUsage verifies bare "/btw" with no prior side thread does not +// launch a run and prints usage guidance. +func TestBtwBareUsage(t *testing.T) { + p := &replProvider{reply: "unused"} + deps, _ := newTestDeps(t, p) + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/btw\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Fatalf("bare /btw must not launch a run, got %d calls", p.calls) + } + if !strings.Contains(out.String(), "usage: /btw") { + t.Errorf("expected usage hint for bare /btw, got: %q", out.String()) + } +} + +// TestBtwBareReopensLastThread verifies that after a side thread exists, a bare +// "/btw" reopens it, replays the prior side Q&A, and lets the user keep asking +// in the SAME thread. The main context stays untouched throughout. +func TestBtwBareReopensLastThread(t *testing.T) { + p := &replProvider{reply: "side answer"} + deps, _ := newTestDeps(t, p) + + var out bytes.Buffer + // Open a side thread and ask once, leave it, then bare /btw reopens it and + // asks a follow-up, then leave again and exit the REPL. + in := strings.NewReader("/btw first question?\n/exit\n/btw\nsecond question?\n/exit\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 2 { + t.Fatalf("expected 2 side runs (1 initial + 1 after reopen), got %d", p.calls) + } + if len(deps.agentCtx.Messages) != 0 { + t.Fatalf("main context must stay untouched, got %d messages", len(deps.agentCtx.Messages)) + } + s := out.String() + // The reopen must not print the bare-/btw usage hint (a thread existed). + if strings.Contains(s, "usage: /btw") { + t.Errorf("bare /btw with an existing thread must not print usage, got: %q", s) + } + // The replay must echo the earlier question. + if !strings.Contains(s, "first question?") { + t.Errorf("reopen should replay the earlier side question, got: %q", s) + } +} + +// TestBtwFollowUpsShareThread verifies that after "/btw " the user can ask +// follow-ups at the btw prompt (without retyping /btw), each launching a run, +// and that none of them pollute the main context. "/exit" leaves the thread. +func TestBtwFollowUpsShareThread(t *testing.T) { + p := &replProvider{reply: "ok"} + deps, _ := newTestDeps(t, p) + + var out bytes.Buffer + // First /btw asks once; then two bare follow-ups; then /exit leaves the side + // thread; then /exit ends the REPL. + in := strings.NewReader("/btw first?\nsecond?\nthird?\n/exit\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 3 { + t.Fatalf("expected 3 side runs (1 initial + 2 follow-ups), got %d", p.calls) + } + if len(deps.agentCtx.Messages) != 0 { + t.Fatalf("follow-ups must not pollute main context, got %d messages", len(deps.agentCtx.Messages)) + } + if !strings.Contains(out.String(), "left side thread") { + t.Errorf("expected 'left side thread' on /exit from the side thread") + } +} + +// TestBtwFollowUpLoopAccumulates verifies the side context grows across +// follow-ups so a later question sees the earlier Q&A. +func TestBtwFollowUpLoopAccumulates(t *testing.T) { + side := &agentcore.AgentContext{} + deps, _ := newTestDeps(t, &replProvider{reply: "a"}) + setCancel := func(context.CancelFunc) {} + settings := btw.ResolveBtwSettings(&bytes.Buffer{}, &deps) + btw.AskSide(setCancel, &bytes.Buffer{}, &deps, side, settings, "q1") + n1 := len(side.Messages) + btw.AskSide(setCancel, &bytes.Buffer{}, &deps, side, settings, "q2") + if len(side.Messages) <= n1 { + t.Fatalf("side context should accumulate across follow-ups: %d then %d", n1, len(side.Messages)) + } +} + +// appending to the side thread cannot reach the main slice. +func TestNewSideContextIsolated(t *testing.T) { + main := &agentcore.AgentContext{ + SystemPrompt: "sys", + Messages: agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}, + }, + } + side := btw.NewSideContext(main) + if side.SystemPrompt != "sys" { + t.Errorf("side thread should inherit the system prompt") + } + if len(side.Messages) != 1 { + t.Fatalf("side thread should be seeded with the main messages, got %d", len(side.Messages)) + } + side.Messages = append(side.Messages, agentcore.UserMessage{RoleField: agentcore.RoleUser}) + if len(main.Messages) != 1 { + t.Fatalf("appending to the side thread must not grow the main context, got %d", len(main.Messages)) + } +} diff --git a/pigo/internal/cli/repl/color_test.go b/pigo/internal/cli/repl/color_test.go new file mode 100644 index 0000000..2462ed4 --- /dev/null +++ b/pigo/internal/cli/repl/color_test.go @@ -0,0 +1,41 @@ +package repl + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/runtime" +) + +// runtimeHelpRegistry builds a slash registry with the live built-in commands +// (/help, /model, /exit, /quit, …) registered against a throwaway live config, +// mirroring what buildSlashRegistry wires up for the real REPL. +func runtimeHelpRegistry(t *testing.T) *runtime.SlashRegistry { + t.Helper() + reg := runtime.NewSlashRegistry() + prompts.RegisterLiveCommands(reg, &cli.LiveConfig{Model: "faux", ProviderName: "faux"}) + return reg +} + +// TestHelpListingColorized verifies the /help action emits ANSI codes when +// color is enabled — the command names are highlighted, not plain. +func TestHelpListingColorized(t *testing.T) { + t.Setenv("NO_COLOR", "1") // force the deterministic (plain) branch + reg := runtimeHelpRegistry(t) + out, err := reg.ResolveOutcome("/help") + if err != nil { + t.Fatalf("resolve /help: %v", err) + } + // With NO_COLOR the listing must be plain text (no escape codes) and still + // contain the command names. + if strings.Contains(out.Message, "\033[") { + t.Errorf("NO_COLOR listing should carry no escape codes, got %q", out.Message) + } + for _, want := range []string{"/help", "/exit", "/quit"} { + if !strings.Contains(out.Message, want) { + t.Errorf("/help listing missing %q, out=%q", want, out.Message) + } + } +} diff --git a/pigo/internal/cli/repl/dream_repl.go b/pigo/internal/cli/repl/dream_repl.go new file mode 100644 index 0000000..cc7a7b5 --- /dev/null +++ b/pigo/internal/cli/repl/dream_repl.go @@ -0,0 +1,240 @@ +// This file implements the manual `/dream` REPL command (SPEC §4.1, US-007): +// it spawns the process-isolated memory-consolidation subprocess +// (`pigo --dream [--dream-dry-run] -C `), captures the single-line +// Report JSON the child writes to stdout (SPEC §4.2), and renders it as a +// full-table change report. `--dry-run` runs the same analysis without writing +// (the subprocess enforces that; the command only reflects report.DryRun). +// +// The report renderers (RenderReportTable / RenderReportLine) live here rather +// than in internal/dream/report.go so the presentation layer stays in the CLI +// package and to keep dream internals conflict-free. RenderReportLine is +// exported because the startup background trigger (#526) reuses it for its +// one-line summary. +package repl + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "time" + + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/dream" +) + +// dreamSubprocessResult is the parsed outcome of one spawn: the decoded Report +// plus whatever the child wrote to stderr (surfaced in error messages so a +// failing run is diagnosable). It is returned by the spawn seam so the +// parse-and-render path can be unit-tested against canned JSON without spawning +// a real LLM-backed dream. +type dreamSubprocessResult struct { + report dream.Report + stderr string +} + +// spawnDream is the seam that launches the dream subprocess and decodes its +// stdout Report. It is a package var so tests can substitute a canned stdout +// (avoiding a real LLM-backed run — see acceptance criteria). The production +// implementation is spawnDreamSubprocess. +var spawnDream = spawnDreamSubprocess + +// spawnDreamSubprocess runs `pigo --dream [--dream-dry-run] -C ` to +// completion, capturing stdout (the single-line Report JSON, SPEC §4.2) and +// stderr (progress/diagnostics). A non-zero exit or unparseable stdout is +// returned as an error carrying the stderr tail so the caller can print a clear +// failure (SPEC §6.1). It never mutates the REPL's own state. +func spawnDreamSubprocess(ctx context.Context, projectDir string, dryRun bool) (dreamSubprocessResult, error) { + exe, err := os.Executable() + if err != nil { + return dreamSubprocessResult{}, fmt.Errorf("resolve pigo executable: %w", err) + } + args := []string{"--dream"} + if dryRun { + args = append(args, "--dream-dry-run") + } + if projectDir != "" { + args = append(args, "-C", projectDir) + } + var stdout, stderr bytes.Buffer + cmd := exec.CommandContext(ctx, exe, args...) + cmd.Stdout = &stdout + cmd.Stderr = &stderr + runErr := cmd.Run() + errTail := strings.TrimSpace(stderr.String()) + if runErr != nil { + // Exit 1 (or a kill/timeout) → failed (SPEC §6.1). Surface the stderr + // tail so the user sees why. + if errTail != "" { + return dreamSubprocessResult{stderr: errTail}, fmt.Errorf("%w: %s", runErr, errTail) + } + return dreamSubprocessResult{}, runErr + } + var report dream.Report + if err := json.Unmarshal(bytes.TrimSpace(stdout.Bytes()), &report); err != nil { + // Unparseable stdout is treated as failed (SPEC §6.1). + return dreamSubprocessResult{stderr: errTail}, fmt.Errorf("parse report: %w", err) + } + return dreamSubprocessResult{report: report, stderr: errTail}, nil +} + +// runDream handles an intercepted `/dream` (or `/dream --dry-run`) line: it +// checks for a live lock (so a background dream already running yields the +// "已有 dream 在运行" notice rather than a confusing empty report), spawns the +// consolidation subprocess with a progress indication, and renders the returned +// Report as a full table. A failed or unparseable run prints a clear error and +// returns to the prompt without crashing the REPL (SPEC §6.1). +func runDream(out io.Writer, deps replDeps, line string) { + dryRun := dreamHasDryRun(line) + + // Pre-spawn lock check: if a background (or other) dream already holds a live + // lock, the subprocess would just skip and emit an all-zero report, + // indistinguishable from "nothing changed". Detect it here so the manual + // command can tell the user (SPEC §6.1 locked row: "已有 dream 在运行"). + if dreamLockHeld(deps.memoryRoot) { + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Yellow, "已有 dream 在运行 (a dream consolidation is already running)")) + return + } + + progress := "Dreaming… (consolidating memory)" + if dryRun { + progress = "Dreaming… (dry-run, analyzing memory — nothing will be written)" + } + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, progress)) + + // Bound the subprocess so a hung LLM-backed run cannot wedge the REPL + // indefinitely (SPEC §6.3/§11.2: parent context timeout, default 10min). On + // timeout CommandContext kills the child and spawnDream surfaces a failed run. + ctx, cancel := context.WithTimeout(context.Background(), dreamRunTimeout) + defer cancel() + res, err := spawnDream(ctx, deps.cwd, dryRun) + if err != nil { + fmt.Fprintf(out, "%s %v\n", ui.Colorize(ui.Enabled(), ui.Red, "dream failed:"), err) + return + } + RenderReportTable(out, res.report) +} + +// dreamRunTimeout bounds a manual /dream subprocess (SPEC §6.3 default 10min). +var dreamRunTimeout = 10 * time.Minute + +// dreamHasDryRun reports whether the /dream command line carries the --dry-run +// flag. It accepts "--dry-run" as a standalone token so "/dream --dry-run" and +// "/dream --dry-run" both match, while a bare "/dream" does not. +func dreamHasDryRun(line string) bool { + for _, f := range strings.Fields(line) { + if f == "--dry-run" { + return true + } + } + return false +} + +// dreamLockPayload mirrors the on-disk dream.lock body (SPEC §3.1: +// {"pid":..,"started_at":..}). It is decoded read-only in the parent to detect +// a running dream; the authoritative lock logic lives in internal/dream/lock.go. +type dreamLockPayload struct { + StartedAt time.Time `json:"started_at"` +} + +// dreamLockHeld reports whether a live (non-stale) dream lock exists under +// memoryRoot. A missing root/lock or a stale lock (older than +// dream.DefaultStaleAfter, matching the runner's takeover rule) reads as not +// held. Read-only: it never creates, removes, or takes over the lock — that is +// the subprocess Runner's job. +func dreamLockHeld(memoryRoot string) bool { + if memoryRoot == "" { + return false + } + path := filepath.Join(memoryRoot, "global", "dream", "dream.lock") + data, err := os.ReadFile(path) + if err != nil { + return false + } + var info dreamLockPayload + if err := json.Unmarshal(data, &info); err != nil { + // Malformed lock body: the runner treats it as stale/takeable, so it is + // not a live lock from the user's perspective. + return false + } + if info.StartedAt.IsZero() { + return false + } + return time.Since(info.StartedAt) <= dream.DefaultStaleAfter +} + +// RenderReportTable writes the full change report (SPEC §2.2/§6.1 manual row): +// one aligned row per counter, byte/file before→after, and any Notes. A dry-run +// report is clearly labeled DRY-RUN and states that nothing was written. +func RenderReportTable(out io.Writer, r dream.Report) { + enabled := ui.Enabled() + if r.DryRun { + fmt.Fprintln(out, ui.Colorize(enabled, ui.Bold, "dream report [DRY-RUN — nothing written]")) + } else { + fmt.Fprintln(out, ui.Colorize(enabled, ui.Bold, "dream report")) + } + rows := []struct { + label string + value string + }{ + {"merged", fmt.Sprintf("%d", r.Merged)}, + {"deduped", fmt.Sprintf("%d", r.Deduped)}, + {"paths-cleaned", fmt.Sprintf("%d", r.PathsCleaned)}, + {"pruned", fmt.Sprintf("%d", r.Pruned)}, + {"distilled", fmt.Sprintf("%d", r.Distilled)}, + {"bytes", fmt.Sprintf("%s → %s", formatBytes(r.BytesBefore), formatBytes(r.BytesAfter))}, + {"files", fmt.Sprintf("%d → %d", r.FilesBefore, r.FilesAfter)}, + {"reconciled", fmt.Sprintf("indexed %d, pruned %d", r.Reconciled.Indexed, r.Reconciled.Pruned)}, + } + width := 0 + for _, row := range rows { + if len(row.label) > width { + width = len(row.label) + } + } + for _, row := range rows { + label := ui.Colorize(enabled, ui.Dim, fmt.Sprintf(" %-*s", width, row.label)) + fmt.Fprintf(out, "%s %s\n", label, row.value) + } + if len(r.Notes) > 0 { + fmt.Fprintln(out, ui.Colorize(enabled, ui.Dim, " notes:")) + for _, n := range r.Notes { + fmt.Fprintf(out, " - %s\n", n) + } + } +} + +// RenderReportLine renders a compact one-line summary of a dream Report. It is +// exported so the startup background trigger (#526) can reuse it for its +// non-intrusive one-line notice (SPEC §6.1 background row). A dry-run report is +// prefixed [DRY-RUN]. +func RenderReportLine(r dream.Report) string { + prefix := "dream:" + if r.DryRun { + prefix = "dream [DRY-RUN]:" + } + return fmt.Sprintf("%s merged %d, deduped %d, paths-cleaned %d, pruned %d, distilled %d, %s→%s, %d→%d files", + prefix, r.Merged, r.Deduped, r.PathsCleaned, r.Pruned, r.Distilled, + formatBytes(r.BytesBefore), formatBytes(r.BytesAfter), r.FilesBefore, r.FilesAfter) +} + +// formatBytes renders a byte count as a compact human-readable string (B/KB/MB). +// It uses 1024-based units and one decimal place above 1KB, matching the terse +// style of the rest of the REPL status output. +func formatBytes(n int64) string { + const unit = 1024 + if n < unit { + return fmt.Sprintf("%dB", n) + } + div, exp := int64(unit), 0 + for x := n / unit; x >= unit; x /= unit { + div *= unit + exp++ + } + return fmt.Sprintf("%.1f%cB", float64(n)/float64(div), "KMGT"[exp]) +} diff --git a/pigo/internal/cli/repl/dream_repl_test.go b/pigo/internal/cli/repl/dream_repl_test.go new file mode 100644 index 0000000..688db87 --- /dev/null +++ b/pigo/internal/cli/repl/dream_repl_test.go @@ -0,0 +1,255 @@ +package repl + +import ( + "bytes" + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/dream" +) + +// sampleReport is a non-trivial Report used across the renderer tests: every +// counter is distinct so a mis-wired field is caught, and Notes/Reconciled are +// populated so their rendering is exercised too. +func sampleReport() dream.Report { + r := dream.Report{ + Merged: 2, + Deduped: 1, + PathsCleaned: 3, + Pruned: 4, + Distilled: 5, + BytesBefore: 4096, + BytesAfter: 2048, + FilesBefore: 9, + FilesAfter: 7, + Notes: []string{"pruned stale entry X", "no new sessions"}, + } + r.Reconciled.Indexed = 6 + r.Reconciled.Pruned = 1 + return r +} + +func TestRenderReportTable(t *testing.T) { + var buf bytes.Buffer + RenderReportTable(&buf, sampleReport()) + got := buf.String() + + // Key counts must appear (label + value). + for _, want := range []string{ + "dream report", + "merged", "deduped", "paths-cleaned", "pruned", "distilled", + "4.0KB → 2.0KB", // bytes before→after + "9 → 7", // files before→after + "indexed 6, pruned 1", + "pruned stale entry X", + "no new sessions", + } { + if !strings.Contains(got, want) { + t.Errorf("full table missing %q\n---\n%s", want, got) + } + } + // A non-dry-run report must NOT carry the DRY-RUN label. + if strings.Contains(got, "DRY-RUN") { + t.Errorf("non-dry-run table should not show DRY-RUN label:\n%s", got) + } +} + +func TestRenderReportTableDryRun(t *testing.T) { + r := sampleReport() + r.DryRun = true + var buf bytes.Buffer + RenderReportTable(&buf, r) + got := buf.String() + if !strings.Contains(got, "DRY-RUN") { + t.Errorf("dry-run table must show DRY-RUN label:\n%s", got) + } + if !strings.Contains(got, "nothing written") { + t.Errorf("dry-run table should state nothing was written:\n%s", got) + } +} + +func TestRenderReportLine(t *testing.T) { + got := RenderReportLine(sampleReport()) + for _, want := range []string{ + "dream:", + "merged 2", "deduped 1", "paths-cleaned 3", "pruned 4", "distilled 5", + "4.0KB→2.0KB", "9→7 files", + } { + if !strings.Contains(got, want) { + t.Errorf("one-line summary missing %q: %q", want, got) + } + } + if strings.Contains(got, "DRY-RUN") { + t.Errorf("non-dry-run line should not show DRY-RUN: %q", got) + } +} + +func TestRenderReportLineDryRun(t *testing.T) { + r := sampleReport() + r.DryRun = true + got := RenderReportLine(r) + if !strings.Contains(got, "DRY-RUN") { + t.Errorf("dry-run line must show DRY-RUN: %q", got) + } +} + +func TestRenderReportZeroValue(t *testing.T) { + // The zero value is a valid "nothing changed" report and must render cleanly. + var buf bytes.Buffer + RenderReportTable(&buf, dream.Report{}) + if line := RenderReportLine(dream.Report{}); !strings.Contains(line, "merged 0") { + t.Errorf("zero-value line should render zero counts: %q", line) + } + if !strings.Contains(buf.String(), "0B → 0B") { + t.Errorf("zero-value table should render 0B → 0B:\n%s", buf.String()) + } +} + +func TestFormatBytes(t *testing.T) { + cases := []struct { + in int64 + want string + }{ + {0, "0B"}, + {512, "512B"}, + {1024, "1.0KB"}, + {1536, "1.5KB"}, + {1048576, "1.0MB"}, + } + for _, c := range cases { + if got := formatBytes(c.in); got != c.want { + t.Errorf("formatBytes(%d) = %q, want %q", c.in, got, c.want) + } + } +} + +func TestDreamHasDryRun(t *testing.T) { + cases := []struct { + line string + want bool + }{ + {"/dream", false}, + {"/dream --dry-run", true}, + {"/dream --dry-run", true}, + {"/dream --dryrun", false}, + {"/dream extra --dry-run", true}, + } + for _, c := range cases { + if got := dreamHasDryRun(c.line); got != c.want { + t.Errorf("dreamHasDryRun(%q) = %v, want %v", c.line, got, c.want) + } + } +} + +// TestRunDreamRendersCannedReport exercises the parse+render path via the spawn +// seam with a canned Report, without spawning a real LLM-backed dream. +func TestRunDreamRendersCannedReport(t *testing.T) { + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(_ context.Context, _ string, dryRun bool) (dreamSubprocessResult, error) { + r := sampleReport() + r.DryRun = dryRun + return dreamSubprocessResult{report: r}, nil + } + + var buf bytes.Buffer + runDream(&buf, replDeps{}, "/dream") + got := buf.String() + if !strings.Contains(got, "dream report") || !strings.Contains(got, "merged") { + t.Errorf("runDream should render the full table:\n%s", got) + } + if strings.Contains(got, "DRY-RUN") { + t.Errorf("non-dry-run runDream should not show DRY-RUN:\n%s", got) + } +} + +func TestRunDreamDryRunLabel(t *testing.T) { + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(_ context.Context, _ string, dryRun bool) (dreamSubprocessResult, error) { + r := sampleReport() + r.DryRun = dryRun + return dreamSubprocessResult{report: r}, nil + } + var buf bytes.Buffer + runDream(&buf, replDeps{}, "/dream --dry-run") + if !strings.Contains(buf.String(), "DRY-RUN") { + t.Errorf("/dream --dry-run should render DRY-RUN label:\n%s", buf.String()) + } +} + +// TestRunDreamFailure asserts a subprocess failure (exit 1 / unparseable stdout) +// prints a clear error and does not crash the REPL (SPEC §6.1). +func TestRunDreamFailure(t *testing.T) { + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(_ context.Context, _ string, _ bool) (dreamSubprocessResult, error) { + return dreamSubprocessResult{}, errFake + } + var buf bytes.Buffer + runDream(&buf, replDeps{}, "/dream") + if !strings.Contains(buf.String(), "dream failed") { + t.Errorf("failed run should print an error:\n%s", buf.String()) + } +} + +var errFake = &fakeErr{} + +type fakeErr struct{} + +func (*fakeErr) Error() string { return "boom" } + +func TestDreamLockHeld(t *testing.T) { + // No memory root → never held. + if dreamLockHeld("") { + t.Fatal("empty memoryRoot must read as not held") + } + root := t.TempDir() + // No lock file yet → not held. + if dreamLockHeld(root) { + t.Fatal("missing lock must read as not held") + } + // Acquire a real lock via the dream package → held. + lock, err := dream.AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock: %v", err) + } + if !dreamLockHeld(root) { + t.Error("a freshly acquired lock must read as held") + } + // Release → not held. + if err := lock.Release(); err != nil { + t.Fatalf("Release: %v", err) + } + if dreamLockHeld(root) { + t.Error("a released lock must read as not held") + } +} + +// TestRunDreamLockedNotice asserts the manual command surfaces the locked +// message and does NOT spawn when a live lock is present (SPEC §6.1). +func TestRunDreamLockedNotice(t *testing.T) { + root := t.TempDir() + lock, err := dream.AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock: %v", err) + } + t.Cleanup(func() { _ = lock.Release() }) + + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawned := false + spawnDream = func(_ context.Context, _ string, _ bool) (dreamSubprocessResult, error) { + spawned = true + return dreamSubprocessResult{}, nil + } + var buf bytes.Buffer + runDream(&buf, replDeps{memoryRoot: root}, "/dream") + if spawned { + t.Error("runDream must not spawn while a live lock is held") + } + if !strings.Contains(buf.String(), "已有 dream 在运行") { + t.Errorf("locked run should print the locked notice:\n%s", buf.String()) + } +} diff --git a/pigo/internal/cli/repl/dream_startup.go b/pigo/internal/cli/repl/dream_startup.go new file mode 100644 index 0000000..3efb1f9 --- /dev/null +++ b/pigo/internal/cli/repl/dream_startup.go @@ -0,0 +1,58 @@ +// This file wires the startup background trigger for /dream memory +// consolidation (US-008, FR-4/FR-17): when an interactive REPL session starts, +// if [dream].enabled and a consolidation is due (per dream.State.Due), the dream +// subprocess is spawned in the BACKGROUND so it never delays the first user +// response, and a non-intrusive one-line summary is printed on completion. +// +// The decision + goroutine live in dream.Scheduler; this file only supplies the +// CLI-side spawn seam (reusing spawnDream from dream_repl.go) and the one-line +// notice renderer (RenderReportLine). The subprocess's O_EXCL lock enforces +// single-instance, so a second trigger just yields a skipped child (silent). +package repl + +import ( + "context" + "fmt" + "io" + + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/dream" +) + +// dreamStartupScheduler owns the startup auto-trigger decision. It is stateless +// (dream.Scheduler is a zero-size type), so a package value is enough; tests +// exercise this path through the spawnDream seam rather than replacing it. +var dreamStartupScheduler dream.Scheduler + +// maybeStartBackgroundDream launches an auto-consolidation in the background at +// interactive session startup when dream is enabled and due. It never blocks: +// the due check is a single state.json read and any spawn runs in a goroutine, +// so the first prompt is served immediately (SPEC FR-4 / §8.2). The completion +// notice is a single dim line via RenderReportLine; a skipped, no-op, or failed +// run prints nothing (SPEC §6.1). +// +// out is written to from the background goroutine, so callers must pass a writer +// safe for a late async line (the REPL uses os.Stdout, where a one-line notice +// simply appears in scrollback). It is a no-op when memoryRoot is empty (dream +// state has nowhere to live) or dream is disabled/not due. +func maybeStartBackgroundDream(out io.Writer, memoryRoot, projectDir string, cfg dream.Config) bool { + if memoryRoot == "" { + return false + } + return dreamStartupScheduler.MaybeRunBackground(context.Background(), dream.BackgroundDeps{ + MemoryRoot: memoryRoot, + ProjectDir: projectDir, + Config: cfg, + Spawn: func(ctx context.Context, dir string) (dream.Report, error) { + // Bound the background run like a manual /dream (SPEC §6.3): a hung + // LLM-backed pass is killed rather than leaking a goroutine forever. + ctx, cancel := context.WithTimeout(ctx, dreamRunTimeout) + defer cancel() + res, err := spawnDream(ctx, dir, false) + return res.report, err + }, + OnReport: func(r dream.Report) { + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, RenderReportLine(r))) + }, + }) +} diff --git a/pigo/internal/cli/repl/dream_startup_test.go b/pigo/internal/cli/repl/dream_startup_test.go new file mode 100644 index 0000000..b63a121 --- /dev/null +++ b/pigo/internal/cli/repl/dream_startup_test.go @@ -0,0 +1,132 @@ +package repl + +import ( + "bytes" + "context" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/smallnest/pigo/internal/dream" +) + +// syncWriter is a concurrency-safe writer whose first Write closes done, so a +// test can wait for the async one-line notice and then read it under the same +// lock the background goroutine wrote it under (bytes.Buffer is not safe for +// concurrent use). +type syncWriter struct { + mu sync.Mutex + buf bytes.Buffer + done chan struct{} + once sync.Once +} + +func newSyncWriter() *syncWriter { return &syncWriter{done: make(chan struct{})} } + +func (w *syncWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + n, err := w.buf.Write(p) + w.once.Do(func() { close(w.done) }) + return n, err +} + +func (w *syncWriter) String() string { + w.mu.Lock() + defer w.mu.Unlock() + return w.buf.String() +} + +// seedDueDreamState writes a state.json under memoryRoot old enough that a +// 7-day-interval dream is due, so maybeStartBackgroundDream will spawn. +func seedDueDreamState(t *testing.T, memoryRoot string) { + t.Helper() + if err := dream.SaveState(memoryRoot, dream.State{ + LastRunAt: time.Now().Add(-30 * 24 * time.Hour), + LastStatus: "ok", + }); err != nil { + t.Fatalf("SaveState: %v", err) + } +} + +func TestMaybeStartBackgroundDream_NoticeOnChanges(t *testing.T) { + root := t.TempDir() + seedDueDreamState(t, root) + + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(_ context.Context, dir string, dryRun bool) (dreamSubprocessResult, error) { + if dryRun { + t.Errorf("background trigger must not run in dry-run mode") + } + if dir != "/proj/y" { + t.Errorf("spawn got dir %q, want /proj/y", dir) + } + return dreamSubprocessResult{report: dream.Report{Merged: 3}}, nil + } + + w := newSyncWriter() + if !maybeStartBackgroundDream(w, root, "/proj/y", dream.NewConfig(nil, 7, 20)) { + t.Fatal("due+enabled dream should launch a background run") + } + select { + case <-w.done: + case <-time.After(2 * time.Second): + t.Fatal("one-line notice not written within timeout") + } + if got := w.String(); !strings.Contains(got, "dream:") || !strings.Contains(got, "merged 3") { + t.Fatalf("one-line notice missing/incorrect: %q", got) + } +} + +func TestMaybeStartBackgroundDream_DisabledNoSpawn(t *testing.T) { + root := t.TempDir() + seedDueDreamState(t, root) + + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(context.Context, string, bool) (dreamSubprocessResult, error) { + t.Fatal("disabled dream must not spawn a subprocess") + return dreamSubprocessResult{}, nil + } + + enabledFalse := false + var buf bytes.Buffer + if maybeStartBackgroundDream(&buf, root, "/proj/y", dream.NewConfig(&enabledFalse, 7, 20)) { + t.Fatal("disabled dream must not launch") + } +} + +func TestMaybeStartBackgroundDream_EmptyRootNoSpawn(t *testing.T) { + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(context.Context, string, bool) (dreamSubprocessResult, error) { + t.Fatal("empty memory root must not spawn") + return dreamSubprocessResult{}, nil + } + var buf bytes.Buffer + if maybeStartBackgroundDream(&buf, "", "/proj/y", dream.NewConfig(nil, 7, 20)) { + t.Fatal("empty memory root must not launch") + } +} + +func TestMaybeStartBackgroundDream_NeverRunNoSpawn(t *testing.T) { + // A fresh memory root (no state.json) is "never run" → not due → no spawn. + root := filepath.Join(t.TempDir(), "empty") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatal(err) + } + orig := spawnDream + t.Cleanup(func() { spawnDream = orig }) + spawnDream = func(context.Context, string, bool) (dreamSubprocessResult, error) { + t.Fatal("never-run state must not spawn (first run is manual)") + return dreamSubprocessResult{}, nil + } + var buf bytes.Buffer + if maybeStartBackgroundDream(&buf, root, "/proj/y", dream.NewConfig(nil, 7, 20)) { + t.Fatal("never-run dream must not launch") + } +} diff --git a/pigo/internal/cli/repl/help_line_test.go b/pigo/internal/cli/repl/help_line_test.go new file mode 100644 index 0000000..34c2787 --- /dev/null +++ b/pigo/internal/cli/repl/help_line_test.go @@ -0,0 +1,61 @@ +package repl + +// Tests for /help's template listing (US-011, #341): formatHelpLine renders +// "/name - description (source: )" (hint omitted when +// absent), and the /help Action includes prompt templates with their hint and +// source tier alongside built-ins. + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/runtime" +) + +func TestFormatHelpLine(t *testing.T) { + cases := []struct { + name string + cmd runtime.SlashCommand + want string + }{ + {"hint+desc global", runtime.SlashCommand{Name: "review", ArgumentHint: "", Description: "Review PRs", Tier: runtime.TierGlobal}, "/review - Review PRs (source: global)"}, + {"desc only builtin", runtime.SlashCommand{Name: "help", Description: "list commands", Tier: runtime.TierBuiltin}, "/help - list commands (source: builtin)"}, + {"hint only cli", runtime.SlashCommand{Name: "wr", ArgumentHint: "[instructions]", Tier: runtime.TierCLI}, "/wr [instructions] (source: cli)"}, + {"neither project", runtime.SlashCommand{Name: "deploy", Tier: runtime.TierProject}, "/deploy (source: project)"}, + {"settings tier", runtime.SlashCommand{Name: "audit", Description: "audit changelog", Tier: runtime.TierSettings}, "/audit - audit changelog (source: settings)"}, + } + for _, c := range cases { + if got := formatHelpLine(c.cmd); got != c.want { + t.Errorf("%s: got %q, want %q", c.name, got, c.want) + } + } +} + +// TestHelpActionIncludesTemplateLabelAndTier verifies the /help Action lists a +// prompt template with its argument-hint, description, and source tier. +func TestHelpActionIncludesTemplateLabelAndTier(t *testing.T) { + reg := runtime.NewSlashRegistry() + prompts.RegisterLiveCommands(reg, &cli.LiveConfig{Model: "test", ProviderName: "test"}) + reg.AddUser(runtime.SlashCommand{ + Name: "review", + ArgumentHint: "", + Description: "Review PRs", + Tier: runtime.TierGlobal, + Expand: func(string) string { return "" }, + }) + + out, err := reg.ResolveOutcome("/help") + if err != nil { + t.Fatalf("ResolveOutcome /help: %v", err) + } + if !out.Handled || out.Kind != runtime.SlashAction { + t.Fatalf("/help should be a handled action, got handled=%v kind=%v", out.Handled, out.Kind) + } + for _, want := range []string{"/review", "", "Review PRs", "(source: global)"} { + if !strings.Contains(out.Message, want) { + t.Errorf("/help output missing %q:\n%s", want, out.Message) + } + } +} diff --git a/pigo/internal/cli/repl/host.go b/pigo/internal/cli/repl/host.go new file mode 100644 index 0000000..fca6bf8 --- /dev/null +++ b/pigo/internal/cli/repl/host.go @@ -0,0 +1,53 @@ +// This file makes replDeps satisfy cli.Host: the accessor and mutator methods +// let the /goal, /btw, /status and REPL logic reach the session's collaborators +// and mutable state through the cli.Host contract rather than the concrete +// aggregate. The compile-time assertion below fails the build if replDeps drifts +// out of conformance. +package repl + +import ( + "bufio" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +var _ cli.Host = (*replDeps)(nil) + +func (d *replDeps) Store() *session.Store { return d.store } +func (d *replDeps) Header() session.SessionHeader { return d.header } +func (d *replDeps) AgentCtx() *agentcore.AgentContext { return d.agentCtx } +func (d *replDeps) Live() *cli.LiveConfig { return d.live } +func (d *replDeps) Registry() *agenttool.ToolRegistry { return d.reg } +func (d *replDeps) Reminders() *runtime.ReminderRegistry { return d.reminders } +func (d *replDeps) Slash() *runtime.SlashRegistry { return d.slash } +func (d *replDeps) Creds() *provider.CredentialStore { return d.creds } +func (d *replDeps) Notifier() *plugin.EventNotifier { return d.notifier } +func (d *replDeps) NotifierHandle() func(agentcore.AgentEvent) { return d.notifierHandle() } +func (d *replDeps) Trust() *trust.Manager { return d.trust } +func (d *replDeps) Goal() *agenttool.GoalState { return d.goal } +func (d *replDeps) Telemetry() *cli.TelemetryHolder { return d.telemetry } +func (d *replDeps) Dispatcher() *hooks.Dispatcher { return d.dispatcher } +func (d *replDeps) HookDeps() run.HookDeps { return d.hookDeps } +func (d *replDeps) Cwd() string { return d.cwd } +func (d *replDeps) Input() *bufio.Reader { return d.in } +func (d *replDeps) ConfirmMu() *sync.Mutex { return d.confirmMu } + +func (d *replDeps) CurLeaf() string { return d.curLeaf } +func (d *replDeps) SetCurLeaf(id string) { d.curLeaf = id } +func (d *replDeps) Persisted() int { return d.persisted } +func (d *replDeps) SetPersisted(n int) { d.persisted = n } + +func (d *replDeps) LastBtw() *agentcore.AgentContext { return d.lastBtw } +func (d *replDeps) SetLastBtw(ctx *agentcore.AgentContext) { d.lastBtw = ctx } +func (d *replDeps) LastBtwBase() int { return d.lastBtwBase } +func (d *replDeps) SetLastBtwBase(n int) { d.lastBtwBase = n } diff --git a/pigo/internal/cli/repl/interactive.go b/pigo/internal/cli/repl/interactive.go new file mode 100644 index 0000000..9ee24c1 --- /dev/null +++ b/pigo/internal/cli/repl/interactive.go @@ -0,0 +1,261 @@ +// This file wires the line-based REPL (US-003) and session persistence +// (US-024, #43) into the pigo command. When invoked without a prompt on a +// terminal, pigo starts the REPL loop (see repl.go); each run's messages are +// persisted to a local JSONL session so the conversation can be listed, resumed +// and replayed later. +package repl + +import ( + "bufio" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/headless" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/dream" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +// sessionStore returns the session store for the interactive REPL. It is a thin +// alias for headless.SessionStore so the REPL and headless runs share one store +// rooted at ~/.pigo/sessions (or PIGO_HOME). +func sessionStore() (*session.Store, error) { + return headless.SessionStore() +} + +// Options carries the resolved run configuration plus optional +// resume state into Run. +type Options struct { + Model string + ProviderName string + Provider provider.Provider + BaseURL string + APIKey string + Protocol string + // ThinkingLevel is the resolved reasoning-effort level (US-023): it seeds the + // live run config so every REPL turn requests it, until a control command + // changes it. + ThinkingLevel agentcore.ThinkingLevel + Tools []agentcore.AgentTool + SysPrompt string + + // ResumeID, when non-empty, resumes an existing session: its messages seed + // the context and replayed transcript. Otherwise a fresh session is created. + ResumeID string + + // Approve, when true, grants the launch directory session trust before the + // run so the first-launch trust prompt is skipped and side-effect tools run + // without per-call confirmation (mirrors pi's --approve/-a). + Approve bool + // Skills is the pre-loaded skill set (loaded once by setupAgentEnv, shared + // with prompt injection). Each is registered as a /skill-name command. Empty + // under --no-skills, so nothing is registered. + Skills []*runtime.Skill + + // Plugins holds the loaded plugin manager so the REPL can deliver lifecycle + // events to subscribed plugins (US-017, #133). It may be nil (no plugins). + Plugins *plugin.Manager + + // ConfigPrompts holds prompt-template paths from the config.toml `prompts` + // array (settings tier); each is a file or dir loaded non-recursively. + ConfigPrompts []string + // CliPrompts holds --prompt-template paths (CLI tier, repeatable). + CliPrompts []string + // NoPromptTemplates disables all prompt-template discovery (global, project, + // settings, CLI); built-in slash commands are unaffected. Independent of + // --no-skills. + NoPromptTemplates bool + + // Dream is the resolved [dream] configuration (US-008). Run uses it to decide + // whether to launch the startup background consolidation; a zero value + // (Enabled false) disables the auto-trigger entirely. + Dream dream.Config +} + +// Run starts the line-based REPL over a persisted session. It keeps +// a single growing AgentContext across prompts (so turns share history) and +// saves the session's messages after each run completes (see runREPL/streamRun +// in repl.go). +func Run(opts Options) error { + creds := provider.NewCredentialStore(nil) + creds.SetOverride(opts.ProviderName, opts.APIKey) + reg := run.ToolRegistry(opts.Tools) + + store, err := sessionStore() + if err != nil { + return err + } + + // Resolve the launch directory once (pigo does not cd during a session). It is + // the trust key, the directory side-effect tools are gated against, and — new + // for #526 — the value stamped onto a fresh SessionHeader.Cwd so the session is + // attributed to a project and a later /dream pass can distill it under the + // right scope (mirrors headless.headlessCwd). An unresolvable cwd yields "" + // (the session stays unattributed) rather than aborting the session. + cwd, cwdErr := os.Getwd() + + // Establish the session: resume an existing one or create a fresh header. + now := time.Now().UTC() + var ( + agentCtx *agentcore.AgentContext + header session.SessionHeader + history []agentcore.AgentMessage + curLeaf string // active leaf id on resume; "" for a fresh session + ) + if opts.ResumeID != "" { + // Interactive resume always appends a fresh user message before running, + // so a session that ended normally (trailing assistant reply) is resumable + // here. Load the raw session and rebuild the context directly. + h, entries, err := store.LoadEntries(opts.ResumeID) + if err != nil { + return err + } + msgs := make(agentcore.MessageList, len(entries)) + for i, e := range entries { + msgs[i] = e.Message + } + if len(entries) > 0 { + curLeaf = entries[len(entries)-1].ID + } + header = h + agentCtx = &agentcore.AgentContext{SystemPrompt: h.SystemPrompt, Messages: msgs, Tools: opts.Tools} + history = msgs + if agentCtx.SystemPrompt == "" { + agentCtx.SystemPrompt = opts.SysPrompt + } + } else { + agentCtx = &agentcore.AgentContext{SystemPrompt: opts.SysPrompt, Tools: opts.Tools} + header = session.SessionHeader{ + ID: session.NewID(now), + CreatedAt: now, + UpdatedAt: now, + Model: opts.Model, + Provider: opts.ProviderName, + SystemPrompt: opts.SysPrompt, + Cwd: cwd, + } + } + + // live holds the run configuration that a control command (e.g. /model) may + // mutate mid-session. streamRun reads it on each prompt so a model switch + // takes effect on the next turn; header is updated so the switch is persisted + // with the session. + live := &cli.LiveConfig{ + Model: opts.Model, + ProviderName: opts.ProviderName, + Provider: opts.Provider, + BaseURL: opts.BaseURL, + Protocol: opts.Protocol, + ThinkingLevel: opts.ThinkingLevel, + ContextWindow: cli.DefaultContextWindow, + } + + // Project trust (US-018, #134): load the persisted trust store for the + // launch directory. A load failure (e.g. a corrupted trust.json) is + // non-fatal: trust is disabled (mgr stays nil) and the REPL still runs - + // the store is surfaced rather than silently overwritten. cwd is captured + // once above since pigo does not cd during a session; if it cannot be resolved + // trust is disabled too, since an empty cwd would silently never match. + mgr, mgrErr := trust.NewManager(trust.DefaultPath()) + if mgrErr != nil { + fmt.Fprintf(os.Stderr, "pigo: trust store unavailable, trust disabled: %v\n", mgrErr) + mgr = nil + } + if cwdErr != nil && mgr != nil { + fmt.Fprintf(os.Stderr, "pigo: cannot resolve working directory, trust disabled: %v\n", cwdErr) + mgr = nil + } + // in is the shared input reader for the main loop and the tool-call + // confirmation prompt (see repl.go). Wrapping os.Stdin once here means both + // read from the same buffer. + reader := bufio.NewReaderSize(os.Stdin, replScanBufInit) + + // Wire slash-commands: built-ins (compile-time) plus any user templates under + // ~/.pigo/commands (mirrors the commands/*.md convention) plus skills under + // ~/.agents/skills. A load error is non-fatal — the REPL still runs with the + // built-ins. Instance built-ins that need live state (/model, /help) are + // registered against `live`. + slash, err := prompts.BuildSlashRegistry(live, opts.Skills, opts.Plugins, prompts.PromptTemplateSources{ + Settings: opts.ConfigPrompts, + CLI: opts.CliPrompts, + Disable: opts.NoPromptTemplates, + ProjectDir: filepath.Join(cwd, ".pigo", "prompts"), + ProjectTrusted: mgr != nil && mgr.IsTrusted(cwd), + }) + if err != nil { + fmt.Fprintf(os.Stderr, "pigo: slash-commands: %v\n", err) + } + trust.RegisterCommand(slash, mgr, cwd) + + // --approve grants the launch directory session trust up front (mirrors pi's + // --approve/-a), so the first-launch prompt is skipped and side-effect tools + // run without per-call confirmation. Otherwise, on the first launch in an + // undecided directory, ask the user how much to trust it before any tool + // runs. This happens before replay so the trust question is the first thing + // the user sees, not their prior history. + trust.EstablishTrust(os.Stdout, reader, mgr, cwd, opts.Approve) + + // Replay the resumed conversation so the user sees history before re-prompting. + if len(history) > 0 { + replayTranscript(os.Stdout, history) + } + + // Startup background consolidation (US-008, FR-4/FR-17): if dream is enabled + // and due, spawn `pigo --dream` in a goroutine now so it runs while the user + // works — it never blocks the first prompt, and prints a one-line notice on + // completion. The dream state/lock live under dream.ResolveMemoryRoot (the + // same root the subprocess consolidates), independent of whether the memory + // tool is wired into this session. Not-due / disabled is a cheap no-op. + maybeStartBackgroundDream(os.Stdout, dream.ResolveMemoryRoot(), cwd, opts.Dream) + + return runREPL(os.Stdin, os.Stdout, replDeps{ + store: store, + header: header, + agentCtx: agentCtx, + live: live, + reg: reg, + reminders: run.TodoReminders(opts.Tools), + slash: slash, + creds: creds, + trust: mgr, + cwd: cwd, + in: reader, + confirmMu: &sync.Mutex{}, + curLeaf: curLeaf, + persisted: len(history), + memoryRoot: run.MemoryRootFromTools(opts.Tools), + memstore: run.MemoryStoreFromTools(opts.Tools), + snap: run.SnapshotRecorderFromTools(opts.Tools), + jobs: run.BashJobStoreFromTools(opts.Tools), + notifier: plugin.NewEventNotifier(opts.Plugins, os.Stderr), + goal: agenttool.NewGoalState(), + telemetry: cli.NewTelemetryHolder(), + }) +} + +// formatHelpLine renders one slash-command line for /help as +// "/name - description (source: )", omitting the hint +// segment when absent. It is the plain, testable form of the /help line; the +// /help Action applies color on top of the same structure. +func formatHelpLine(c runtime.SlashCommand) string { + s := "/" + c.Name + if c.ArgumentHint != "" { + s += " " + c.ArgumentHint + } + if c.Description != "" { + s += " - " + c.Description + } + s += " (source: " + c.Tier.String() + ")" + return s +} diff --git a/pigo/internal/cli/repl/line_editor.go b/pigo/internal/cli/repl/line_editor.go new file mode 100644 index 0000000..e26358c --- /dev/null +++ b/pigo/internal/cli/repl/line_editor.go @@ -0,0 +1,766 @@ +package repl + +import ( + "bufio" + "fmt" + "io" + "os" + "os/exec" + "sort" + "strconv" + "strings" + "unicode/utf8" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" +) + +// errLineInterrupted aliases cli.ErrLineInterrupted so the editor's own returns +// and tests keep the short local name while subpackages (e.g. /btw) recognize +// the same sentinel through the exported cli.ErrLineInterrupted. +var errLineInterrupted = cli.ErrLineInterrupted + +// mlBuffer models the readLine input as one or more lines with a cursor at +// (row, col), col counted in runes within the current line. A fresh buffer +// holds a single empty line with the cursor at the origin. It is the data +// model the multi-line REPL editor (continuation, Shift+Enter, cross-line +// movement/editing, rendering, history) is built on; single-line input behaves +// exactly as a plain string with the cursor at its end. +type mlBuffer struct { + lines []string + row int + col int +} + +func newMLBuffer() *mlBuffer { return &mlBuffer{lines: []string{""}} } + +// String joins the lines with "\n" for submission. A single empty line yields +// "" and there is never a trailing newline. +func (b *mlBuffer) String() string { return strings.Join(b.lines, "\n") } + +// isEmpty reports whether the buffer is a single empty line. +func (b *mlBuffer) isEmpty() bool { return len(b.lines) == 1 && b.lines[0] == "" } + +// single reports whether the buffer holds exactly one line. +func (b *mlBuffer) single() bool { return len(b.lines) == 1 } + +// line returns the text of the current cursor line. +func (b *mlBuffer) line() string { return b.lines[b.row] } + +// setString replaces the whole buffer with s (which may contain "\n"), placing +// the cursor at the end of the last line. Used when the entire input is swapped +// wholesale — accepting a suggestion, browsing history, or restoring a +// multi-line entry. +func (b *mlBuffer) setString(s string) { + b.lines = strings.Split(s, "\n") + b.row = len(b.lines) - 1 + b.col = utf8.RuneCountInString(b.lines[b.row]) +} + +// insert adds s (which must not contain "\n") at the cursor on the current +// line and advances col past it. +func (b *mlBuffer) insert(s string) { + line := b.lines[b.row] + off := runeOffset(line, b.col) + b.lines[b.row] = line[:off] + s + line[off:] + b.col += utf8.RuneCountInString(s) +} + +// backspace deletes the rune immediately left of the cursor on the current +// line. At column 0 it merges the current line into the previous one (the +// cursor landing at the merge point), unless already on the first line, where +// it is a no-op. Multi-byte runes are removed whole. +func (b *mlBuffer) backspace() { + if b.col == 0 { + if b.row == 0 { + return + } + prev := b.lines[b.row-1] + b.col = utf8.RuneCountInString(prev) + b.lines[b.row-1] = prev + b.lines[b.row] + b.lines = append(b.lines[:b.row], b.lines[b.row+1:]...) + b.row-- + return + } + line := b.lines[b.row] + start := runeOffset(line, b.col-1) + end := runeOffset(line, b.col) + b.lines[b.row] = line[:start] + line[end:] + b.col-- +} + +// newline splits the current line at the cursor, moving the text right of the +// cursor onto a fresh line below and placing the cursor at its start. It is how +// Shift+Enter (and, later, backslash continuation) turn one line into two. +func (b *mlBuffer) newline() { + line := b.lines[b.row] + off := runeOffset(line, b.col) + head, tail := line[:off], line[off:] + rest := append([]string{}, b.lines[b.row+1:]...) + b.lines = append(b.lines[:b.row], head, tail) + b.lines = append(b.lines, rest...) + b.row++ + b.col = 0 +} + +// left moves the cursor one rune left, crossing to the end of the previous +// line when already at column 0. It is a no-op at the buffer origin. +func (b *mlBuffer) left() { + if b.col > 0 { + b.col-- + } else if b.row > 0 { + b.row-- + b.col = utf8.RuneCountInString(b.lines[b.row]) + } +} + +// right moves the cursor one rune right, crossing to the start of the next +// line when already at the end of the current line. It is a no-op at the end +// of the last line. +func (b *mlBuffer) right() { + if b.col < utf8.RuneCountInString(b.lines[b.row]) { + b.col++ + } else if b.row < len(b.lines)-1 { + b.row++ + b.col = 0 + } +} + +// up moves the cursor to the previous line, clamping the column to that line's +// length. It is a no-op on the first line. +func (b *mlBuffer) up() { + if b.row > 0 { + b.row-- + if n := utf8.RuneCountInString(b.lines[b.row]); b.col > n { + b.col = n + } + } +} + +// down moves the cursor to the next line, clamping the column to that line's +// length. It is a no-op on the last line. +func (b *mlBuffer) down() { + if b.row < len(b.lines)-1 { + b.row++ + if n := utf8.RuneCountInString(b.lines[b.row]); b.col > n { + b.col = n + } + } +} + +// home moves the cursor to the start of the current line. +func (b *mlBuffer) home() { b.col = 0 } + +// end moves the cursor to the end of the current line. +func (b *mlBuffer) end() { b.col = utf8.RuneCountInString(b.lines[b.row]) } + +// enterContinues collapses the current line's trailing backslash run and +// reports whether a pressed Enter should continue onto a new line instead of +// submitting. A run of k trailing backslashes pairs up as k/2 literal +// backslashes (each "\\" → one "\"); an odd run has one extra backslash that +// escapes the newline, so the caller inserts a continuation line. The run is +// always collapsed to k/2 backslashes, so the escaping "\" never survives into +// submitted text. +func (b *mlBuffer) enterContinues() bool { + line := b.lines[b.row] + k := 0 + for i := len(line) - 1; i >= 0 && line[i] == '\\'; i-- { + k++ + } + if k == 0 { + return false + } + b.lines[b.row] = line[:len(line)-k] + strings.Repeat("\\", k/2) + if n := utf8.RuneCountInString(b.lines[b.row]); b.col > n { + b.col = n + } + return k%2 == 1 +} + +// visibleWidth returns the number of terminal columns s occupies, skipping ANSI +// CSI escape sequences so a colored prompt still aligns its continuation lines. +// Wide runes (CJK ideographs, fullwidth forms, most emoji) count as two columns. +func visibleWidth(s string) int { + w := 0 + for i := 0; i < len(s); { + if s[i] == 0x1b { + i++ + if i < len(s) && s[i] == '[' { + i++ + for i < len(s) && !(s[i] >= 0x40 && s[i] <= 0x7e) { + i++ + } + } + if i < len(s) { + i++ + } + continue + } + r, size := utf8.DecodeRuneInString(s[i:]) + i += size + w += runeWidth(r) + } + return w +} + +// displayWidth returns the number of terminal columns the plain string s +// occupies, summing each rune's cell width. Unlike visibleWidth it does not +// strip ANSI escapes — callers pass already-plain buffer text. +func displayWidth(s string) int { + w := 0 + for _, r := range s { + w += runeWidth(r) + } + return w +} + +// runeWidth reports how many terminal cells a rune occupies: 0 for combining / +// zero-width marks, 2 for East Asian wide and fullwidth characters (and most +// emoji), 1 otherwise. This is what keeps the cursor aligned when the line +// contains CJK text, where one rune spans two columns. +func runeWidth(r rune) int { + switch { + case r == 0: + return 0 + case (r >= 0x0300 && r <= 0x036F), // combining diacritical marks + (r >= 0x1AB0 && r <= 0x1AFF), // combining diacritical marks extended + (r >= 0x1DC0 && r <= 0x1DFF), // combining diacritical marks supplement + (r >= 0x20D0 && r <= 0x20FF), // combining marks for symbols + (r >= 0xFE20 && r <= 0xFE2F), // combining half marks + r == 0x200B: // zero width space + return 0 + case (r >= 0x1100 && r <= 0x115F), // Hangul Jamo + (r >= 0x2E80 && r <= 0x303E), // CJK radicals, Kangxi, CJK symbols + (r >= 0x3041 && r <= 0x33FF), // Hiragana, Katakana, CJK compat + (r >= 0x3400 && r <= 0x4DBF), // CJK Ext A + (r >= 0x4E00 && r <= 0x9FFF), // CJK Unified Ideographs + (r >= 0xA000 && r <= 0xA4CF), // Yi + (r >= 0xAC00 && r <= 0xD7A3), // Hangul syllables + (r >= 0xF900 && r <= 0xFAFF), // CJK compat ideographs + (r >= 0xFE10 && r <= 0xFE19), // vertical forms + (r >= 0xFE30 && r <= 0xFE6F), // CJK compat forms + (r >= 0xFF00 && r <= 0xFF60), // fullwidth forms + (r >= 0xFFE0 && r <= 0xFFE6), // fullwidth signs + (r >= 0x1F300 && r <= 0x1FAFF), // emoji & pictographs + (r >= 0x20000 && r <= 0x3FFFD): // CJK Ext B and beyond + return 2 + default: + return 1 + } +} + +// runeOffset converts a rune column into a byte offset within s. +func runeOffset(s string, col int) int { + off := 0 + for i := 0; i < col && off < len(s); i++ { + _, size := utf8.DecodeRuneInString(s[off:]) + off += size + } + return off +} + +// replLineEditor adds a small shell-style editing layer without turning the +// line-oriented REPL back into a full-screen TUI. On terminals it shows the +// best completion in dim text as the user types. Pipes and tests keep using the +// ordinary buffered reader. +type replLineEditor struct { + in *bufio.Reader + terminal *os.File + out io.Writer + slash *runtime.SlashRegistry + history []string // oldest to newest + models []string +} + +func newREPLLineEditor(in io.Reader, buffered *bufio.Reader, out io.Writer, slash *runtime.SlashRegistry, history []string) *replLineEditor { + e := &replLineEditor{in: buffered, out: out, slash: slash} + if f, ok := in.(*os.File); ok { + if info, err := f.Stat(); err == nil && info.Mode()&os.ModeCharDevice != 0 { + e.terminal = f + } + } + for _, h := range history { + e.remember(h) + } + seen := map[string]bool{} + for _, m := range provider.PresetCatalog { + if !seen[m.ID] { + e.models = append(e.models, m.ID) + seen[m.ID] = true + } + } + return e +} + +func (e *replLineEditor) remember(line string) { + line = strings.TrimSpace(line) + if line == "" { + return + } + e.history = append(e.history, line) + if len(e.history) > 200 { + e.history = e.history[len(e.history)-200:] + } +} + +// formatSlashAutocompleteLabel renders a slash command for the Tab-completion +// hint as "name - description" (mirrors pi's autocomplete). The +// argument-hint is shown verbatim (frontmatter supplies its own / +// [square] brackets); it and the description are omitted when absent, so a +// bare command renders as just its name. +func formatSlashAutocompleteLabel(cmd runtime.SlashCommand) string { + label := cmd.Name + if cmd.ArgumentHint != "" { + label += " " + cmd.ArgumentHint + } + if cmd.Description != "" { + label += " - " + cmd.Description + } + return label +} + +// suggestion returns the single best completion for input, or "" when there is +// none. It is the head of the ordered candidate list (see suggestions). +func (e *replLineEditor) suggestion(input string) string { + if cands := e.suggestions(input); len(cands) > 0 { + return cands[0] + } + return "" +} + +// suggestions returns every completion candidate for input, best first, so the +// caller can cycle through them with the arrow keys. Candidates are gathered in +// priority order — slash commands, then recent inputs, then the /model catalog +// — deduplicated, with the raw input itself excluded. +func (e *replLineEditor) suggestions(input string) []string { + if input == "" { + return nil + } + lower := strings.ToLower(input) + + var out []string + seen := map[string]bool{} + add := func(s string) { + if s == "" || s == input || seen[s] { + return + } + seen[s] = true + out = append(out, s) + } + + if strings.HasPrefix(input, "/") && !strings.ContainsAny(input, " \t") { + var commands []string + for _, cmd := range e.slash.List() { + commands = append(commands, "/"+cmd.Name) + } + sort.Strings(commands) + for _, cmd := range commands { + if strings.HasPrefix(strings.ToLower(cmd), lower) { + add(cmd) + } + } + } + for i := len(e.history) - 1; i >= 0; i-- { + if strings.HasPrefix(strings.ToLower(e.history[i]), lower) { + add(e.history[i]) + } + } + if strings.HasPrefix(lower, "/model ") { + query := strings.TrimSpace(input[len("/model "):]) + for i := len(e.history) - 1; i >= 0; i-- { + h := e.history[i] + if !strings.HasPrefix(h, "/model ") { + continue + } + id := strings.TrimSpace(h[len("/model "):]) + if query == "" || modelMatches(id, query) { + add(h) + } + } + for _, id := range e.models { + if query == "" || modelMatches(id, query) { + add("/model " + id) + } + } + } + return out +} + +func modelMatches(id, query string) bool { + id, query = strings.ToLower(id), strings.ToLower(query) + if strings.HasPrefix(id, query) { + return true + } + if slash := strings.LastIndexByte(id, '/'); slash >= 0 { + return strings.HasPrefix(id[slash+1:], query) + } + return false +} + +// parseCSIParams splits a CSI-u parameter list ("[;]") into the key +// code and modifier. A missing modifier defaults to 1 (no modifier). +func parseCSIParams(params []byte) (code, mod int) { + parts := strings.Split(string(params), ";") + code = atoiDefault(parts[0], 0) + mod = 1 + if len(parts) > 1 { + mod = atoiDefault(parts[1], 1) + } + return code, mod +} + +func atoiDefault(s string, def int) int { + if n, err := strconv.Atoi(s); err == nil { + return n + } + return def +} + +func (e *replLineEditor) ReadLine(prompt string) (string, error) { + if e.terminal == nil { + fmt.Fprint(e.out, prompt) + return e.in.ReadString('\n') + } + probe := exec.Command("stty", "-g") + probe.Stdin = e.terminal + state, err := probe.CombinedOutput() + if err != nil { + fmt.Fprint(e.out, prompt) + return e.in.ReadString('\n') + } + raw := exec.Command("stty", "raw", "-echo") + raw.Stdin = e.terminal + if err := raw.Run(); err != nil { + fmt.Fprint(e.out, prompt) + return e.in.ReadString('\n') + } + // Ask the terminal to report modified keys so Shift+Enter is distinguishable + // from a bare Enter: enable xterm modifyOtherKeys level 1 and push a CSI-u + // (fixterms/kitty) keyboard mode. Level 1 (not 2) reports only keys without + // a standard encoding — so Shift+Enter is escaped while ordinary Tab/Enter + // stay untouched. Terminals that ignore these simply never send the reports, + // and the user falls back to backslash continuation. + fmt.Fprint(e.out, "\x1b[>4;1m\x1b[>1u") + defer func() { + // Restore the terminal's key reporting before the stty state, so we + // never leave it stuck in CSI-u/modifyOtherKeys mode on any exit path. + fmt.Fprint(e.out, "\x1b[4;0m") + restore := exec.Command("stty", strings.TrimSpace(string(state))) + restore.Stdin = e.terminal + _ = restore.Run() + }() + + return e.editLoop(prompt) +} + +// editLoop runs the raw-mode key-processing loop over e.in, kept separate from +// readLine's terminal setup so it can be driven by a programmable io.Reader in +// tests (no real TTY required). It returns the submitted text (lines joined by +// "\n") or an error. +func (e *replLineEditor) editLoop(prompt string) (string, error) { + // buf models the input as a multi-line buffer with a cursor. For this + // single-line editing layer the cursor stays at the end of the sole line, + // so buf behaves exactly like the former input string; the buffer model is + // what later cross-line editing/rendering is built on. + buf := newMLBuffer() + // selected indexes into the current candidate list. It advances with the + // up/down arrows so the user can cycle through suggestions; it resets to 0 + // (the best match) whenever the input text changes, since the candidate list + // is recomputed from scratch. + selected := 0 + // histNav tracks the position while browsing prior inputs with the arrow + // keys on a blank line: -1 means not browsing, otherwise it indexes into + // e.history (oldest to newest). It resets to -1 whenever the user edits the + // line, so history browsing is only active while stepping through entries. + histNav := -1 + // visible returns the suggestion currently shown/accepted: the candidate at + // the selected index, clamped to the available list. + visible := func() string { + cands := e.suggestions(buf.String()) + if len(cands) == 0 { + return "" + } + if selected >= len(cands) { + selected = len(cands) - 1 + } + if selected < 0 { + selected = 0 + } + return cands[selected] + } + // promptW is the prompt's visible width; continuation lines are indented to + // that column so every line's text starts at the same place, with a dim + // marker standing in for the prompt. + promptW := visibleWidth(prompt) + contPrefix := strings.Repeat(" ", promptW) + if promptW >= 2 { + contPrefix = strings.Repeat(" ", promptW-2) + "\033[2m·\033[0m " + } + // prevCursorRow is the screen row (relative to the block's first line) the + // cursor was left on by the previous render, so the next render can climb + // back to the top of the block before clearing and redrawing it. + prevCursorRow := 0 + render := func() { + // Return to the top-left of the block drawn last time and clear it plus + // anything below, so shrinking the buffer leaves no stale rows/chars. + if prevCursorRow > 0 { + fmt.Fprintf(e.out, "\033[%dA", prevCursorRow) + } + fmt.Fprint(e.out, "\r\033[J") + for i, line := range buf.lines { + if i == 0 { + fmt.Fprintf(e.out, "%s%s", prompt, line) + } else { + fmt.Fprintf(e.out, "\r\n%s%s", contPrefix, line) + } + } + // The dim completion hint only fits on a single line with the cursor at + // its end, where it can't collide with continuation rows. + if buf.single() && buf.col == utf8.RuneCountInString(buf.lines[0]) { + if s := visible(); s != "" { + input := buf.lines[0] + if strings.HasPrefix(s, "/") { + // Slash command: render the argument-hint + description label + // (mirrors pi's autocomplete) instead of the bare name suffix. + label := s + if cmd, ok := e.slash.Lookup(s[1:]); ok { + label = formatSlashAutocompleteLabel(cmd) + } + fmt.Fprintf(e.out, "\033[2m -> %s\033[0m", label) + } else if strings.HasPrefix(s, input) { + fmt.Fprintf(e.out, "\033[2m%s\033[0m", s[len(input):]) + } else { + fmt.Fprintf(e.out, "\033[2m → %s\033[0m", s) + } + } + } + // Reposition to the logical (row, col): after the draw the cursor sits + // at the end of the last line, so climb to the target row, then step + // right past the prefix and the display width of the runes left of the + // cursor (wide CJK runes span two columns, so count cells, not runes). + if up := len(buf.lines) - 1 - buf.row; up > 0 { + fmt.Fprintf(e.out, "\033[%dA", up) + } + fmt.Fprint(e.out, "\r") + curLine := buf.lines[buf.row] + cursorCells := displayWidth(curLine[:runeOffset(curLine, buf.col)]) + if col := promptW + cursorCells; col > 0 { + fmt.Fprintf(e.out, "\033[%dC", col) + } + prevCursorRow = buf.row + } + // tryEnter handles a pressed Enter shared by the raw, CSI-u, and + // modifyOtherKeys report paths: if the current line ends with an unescaped + // backslash it continues onto a new line and reports submitted=false; + // otherwise it emits the newline and reports the block as submitted. + tryEnter := func() (string, bool) { + if buf.enterContinues() { + buf.end() + buf.newline() + selected = 0 + histNav = -1 + return "", false + } + // Move below the whole rendered block before the newline so the + // submitted lines stay on screen and the next prompt starts clean. + if down := len(buf.lines) - 1 - buf.row; down > 0 { + fmt.Fprintf(e.out, "\033[%dB", down) + } + fmt.Fprint(e.out, "\r\n") + return buf.String(), true + } + render() + for { + b, err := e.in.ReadByte() + if err != nil { + return buf.String(), err + } + switch b { + case '\r', '\n': + if res, done := tryEnter(); done { + return res, nil + } + case 1: // Ctrl+A moves to line start. + buf.home() + case 5: // Ctrl+E moves to line end. + buf.end() + case 3: // Ctrl+C + fmt.Fprint(e.out, "^C\r\n") + return "", errLineInterrupted + case 4: // Ctrl+D + if buf.isEmpty() { + fmt.Fprint(e.out, "\r\n") + return "", io.EOF + } + case 9: // Tab accepts the visible suggestion. + if s := visible(); s != "" { + buf.setString(s) + selected = 0 + histNav = -1 + } + case 8, 127: + buf.backspace() + selected = 0 + histNav = -1 + case 27: + // Parse a full CSI sequence so multi-parameter reports (CSI-u key + // events like Shift+Enter's \x1b[13;2u) are handled, not just the + // bare arrow sequences. → accepts the visible suggestion, ↑/↓ cycle + // candidates or browse history on a blank line, and Enter reports + // (code 13) either submit or insert a newline depending on the + // modifier. Any other sequence is consumed and ignored so it never + // leaks into the submitted text. + b2, escErr := e.in.ReadByte() + if escErr != nil { + return buf.String(), escErr + } + if b2 == '[' { + var params []byte + var final byte + for { + c, cErr := e.in.ReadByte() + if cErr != nil { + return buf.String(), cErr + } + if c >= 0x40 && c <= 0x7e { + final = c + break + } + params = append(params, c) + } + switch final { + case 'u': // CSI-u key report: "[;]u". + code, mod := parseCSIParams(params) + ctrl := (mod-1)&4 != 0 + switch { + case code == 13: + if mod >= 2 { + buf.newline() + selected = 0 + histNav = -1 + } else if res, done := tryEnter(); done { + return res, nil + } + case ctrl && code == 'd': + // Ctrl+D: EOF on an empty line. Under the kitty keyboard + // protocol (enabled via \x1b[>1u) the terminal reports it here + // as a CSI-u event, not the raw 0x04 byte the case-4 arm handles. + if buf.isEmpty() { + fmt.Fprint(e.out, "\r\n") + return "", io.EOF + } + case ctrl && code == 'c': + // Ctrl+C: same story — delivered as a CSI-u report rather than + // the raw 0x03 byte once the kitty keyboard mode is active. + fmt.Fprint(e.out, "^C\r\n") + return "", errLineInterrupted + } + case '~': // modifyOtherKeys ("27;;~") or Home/End ("1~"/"4~"). + parts := strings.Split(string(params), ";") + if len(parts) == 3 && atoiDefault(parts[0], -1) == 27 { + mod := atoiDefault(parts[1], 1) + code := atoiDefault(parts[2], 0) + if code == 13 { + if mod >= 2 { + buf.newline() + selected = 0 + histNav = -1 + } else if res, done := tryEnter(); done { + return res, nil + } + } + } else if len(parts) == 1 { + switch atoiDefault(parts[0], -1) { + case 1, 7: // Home + buf.home() + case 4, 8: // End + buf.end() + } + } + case 'C': // right arrow: accept a visible suggestion, else move the cursor + if len(params) == 0 { + if s := visible(); s != "" { + buf.setString(s) + selected = 0 + histNav = -1 + } else { + buf.right() + } + } + case 'D': // left arrow moves the cursor (cross-line at column 0) + if len(params) == 0 { + buf.left() + } + case 'H': // Home + if len(params) == 0 { + buf.home() + } + case 'F': // End + if len(params) == 0 { + buf.end() + } + case 'A': // up arrow + if len(params) == 0 { + if !buf.single() { + buf.up() + } else if buf.isEmpty() || histNav >= 0 { + // Browse history: step toward older entries. + if histNav < 0 { + histNav = len(e.history) + } + if histNav > 0 { + histNav-- + buf.setString(e.history[histNav]) + selected = 0 + } + } else if n := len(e.suggestions(buf.String())); n > 0 { + selected = (selected - 1 + n) % n + } + } + case 'B': // down arrow + if len(params) == 0 { + if !buf.single() { + buf.down() + } else if histNav >= 0 { + // Browse history: step toward newer entries; past the + // newest, return to a blank line. + if histNav < len(e.history)-1 { + histNav++ + buf.setString(e.history[histNav]) + } else { + histNav = -1 + buf.setString("") + } + selected = 0 + } else if n := len(e.suggestions(buf.String())); n > 0 { + selected = (selected + 1) % n + } + } + } + } + default: + bytes := []byte{b} + want := 1 + switch { + case b&0xe0 == 0xc0: + want = 2 + case b&0xf0 == 0xe0: + want = 3 + case b&0xf8 == 0xf0: + want = 4 + } + for len(bytes) < want { + next, readErr := e.in.ReadByte() + if readErr != nil { + return buf.String(), readErr + } + bytes = append(bytes, next) + } + buf.insert(string(bytes)) + selected = 0 + histNav = -1 + } + render() + } +} diff --git a/pigo/internal/cli/repl/line_editor_test.go b/pigo/internal/cli/repl/line_editor_test.go new file mode 100644 index 0000000..441038d --- /dev/null +++ b/pigo/internal/cli/repl/line_editor_test.go @@ -0,0 +1,599 @@ +package repl + +import ( + "bufio" + "bytes" + "errors" + "io" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/runtime" +) + +func testLineEditor(history ...string) *replLineEditor { + reg := runtime.NewSlashRegistry() + reg.AddBuiltin(runtime.SlashCommand{Name: "model", Action: func(string) string { return "" }}) + reg.AddBuiltin(runtime.SlashCommand{Name: "models", Action: func(string) string { return "" }}) + return newREPLLineEditor(strings.NewReader(""), bufio.NewReader(strings.NewReader("")), io.Discard, reg, history) +} + +func TestLineEditorPrefersMostRecentMatchingInput(t *testing.T) { + e := testLineEditor("explain old", "other", "explain recent") + if got := e.suggestion("exp"); got != "explain recent" { + t.Fatalf("suggestion = %q, want most recent match", got) + } +} + +func TestLineEditorCompletesSlashCommands(t *testing.T) { + e := testLineEditor() + if got := e.suggestion("/mod"); got != "/model" { + t.Fatalf("suggestion = %q, want /model", got) + } +} + +func TestLineEditorCompletesModelsByRecentUseAndBasename(t *testing.T) { + e := testLineEditor("/model openai/gpt-4o") + if got := e.suggestion("/model "); got != "/model openai/gpt-4o" { + t.Fatalf("empty model suggestion = %q", got) + } + if got := e.suggestion("/model gpt"); got != "/model openai/gpt-4o" { + t.Fatalf("recent model suggestion = %q", got) + } + e = testLineEditor() + got := e.suggestion("/model deepseek") + if got == "" || !strings.HasPrefix(got, "/model ") { + t.Fatalf("catalog model suggestion = %q", got) + } +} + +func TestLineEditorSuggestionsAreOrderedAndDeduped(t *testing.T) { + // Two recent inputs plus a slash command all sharing a prefix: the caller + // cycles this list with the arrow keys, so ordering (best first) and + // dedup both matter. + e := testLineEditor("explain old", "explain recent", "explain recent") + cands := e.suggestions("exp") + if len(cands) != 2 { + t.Fatalf("suggestions = %v, want 2 unique candidates", cands) + } + if cands[0] != "explain recent" || cands[1] != "explain old" { + t.Fatalf("suggestions = %v, want most-recent first", cands) + } + // The head of the list must match the single-suggestion helper. + if e.suggestion("exp") != cands[0] { + t.Fatalf("suggestion head %q != suggestions[0] %q", e.suggestion("exp"), cands[0]) + } +} + +func TestLineEditorSlashCommandsCycleAllMatches(t *testing.T) { + e := testLineEditor() + cands := e.suggestions("/mode") + // Both /model and /models share the prefix; cycling must expose both. + if len(cands) != 2 || cands[0] != "/model" || cands[1] != "/models" { + t.Fatalf("suggestions = %v, want [/model /models]", cands) + } +} + +func TestMLBufferEmptyBehavesLikeEmptyString(t *testing.T) { + b := newMLBuffer() + if !b.isEmpty() { + t.Fatalf("fresh buffer should be empty") + } + if !b.single() { + t.Fatalf("fresh buffer should be single-line") + } + if got := b.String(); got != "" { + t.Fatalf("empty buffer String() = %q, want \"\"", got) + } +} + +func TestMLBufferSingleLineInsert(t *testing.T) { + b := newMLBuffer() + b.insert("hello") + if got := b.String(); got != "hello" { + t.Fatalf("String() = %q, want %q", got, "hello") + } + if b.col != 5 { + t.Fatalf("col = %d, want 5", b.col) + } + if b.isEmpty() { + t.Fatalf("buffer with text should not be empty") + } +} + +func TestMLBufferBackspaceRemovesLastRune(t *testing.T) { + b := newMLBuffer() + b.insert("héllo") // multi-byte rune to exercise UTF-8 handling + b.backspace() + if got := b.String(); got != "héll" { + t.Fatalf("after backspace String() = %q, want %q", got, "héll") + } + // Remove the multi-byte rune specifically. + b.setString("café") + b.backspace() + if got := b.String(); got != "caf" { + t.Fatalf("multi-byte backspace String() = %q, want %q", got, "caf") + } +} + +func TestMLBufferBackspaceEmptyIsNoop(t *testing.T) { + b := newMLBuffer() + b.backspace() + if got := b.String(); got != "" { + t.Fatalf("backspace on empty buffer String() = %q, want \"\"", got) + } + if b.col != 0 || b.row != 0 { + t.Fatalf("cursor moved on empty backspace: row=%d col=%d", b.row, b.col) + } +} + +func TestMLBufferSetStringJoinsWithNewline(t *testing.T) { + b := newMLBuffer() + b.setString("line1\nline2\nline3") + if got := b.String(); got != "line1\nline2\nline3" { + t.Fatalf("String() = %q, want round-trip", got) + } + if b.single() { + t.Fatalf("multi-line buffer reported single()") + } + // Cursor lands at end of last line. + if b.row != 2 || b.col != 5 { + t.Fatalf("cursor = (%d,%d), want (2,5)", b.row, b.col) + } + // No trailing newline is introduced. + if strings.HasSuffix(b.String(), "\n") { + t.Fatalf("String() has trailing newline: %q", b.String()) + } +} + +// editorWithOutput is editorWithInput but captures the editor's terminal +// output so raw-mode rendering (escape sequences) can be asserted. +func editorWithOutput(input string, out io.Writer, history ...string) *replLineEditor { + reg := runtime.NewSlashRegistry() + reg.AddBuiltin(runtime.SlashCommand{Name: "model", Action: func(string) string { return "" }}) + r := strings.NewReader(input) + return newREPLLineEditor(r, bufio.NewReader(r), out, reg, history) +} + +func TestVisibleWidthSkipsANSI(t *testing.T) { + if w := visibleWidth("pigo> "); w != 6 { + t.Fatalf("plain width = %d, want 6", w) + } + if w := visibleWidth("\033[2m·\033[0m "); w != 2 { + t.Fatalf("dim marker width = %d, want 2 (marker + space)", w) + } + if w := visibleWidth("café> "); w != 6 { + t.Fatalf("multi-byte width = %d, want 6 runes", w) + } +} + +func TestRuneAndDisplayWidthHandleWideRunes(t *testing.T) { + // ASCII and Latin-1 accents are one cell; CJK ideographs and fullwidth + // forms are two; combining marks are zero. + if w := runeWidth('a'); w != 1 { + t.Fatalf("width('a') = %d, want 1", w) + } + if w := runeWidth('é'); w != 1 { + t.Fatalf("width('é') = %d, want 1", w) + } + if w := runeWidth('中'); w != 2 { + t.Fatalf("width('中') = %d, want 2", w) + } + if w := runeWidth('́'); w != 0 { // combining acute accent + t.Fatalf("width(combining) = %d, want 0", w) + } + // A mixed CJK/ASCII string sums per-cell. + if w := displayWidth("你好a"); w != 5 { + t.Fatalf("displayWidth(\"你好a\") = %d, want 5", w) + } + // visibleWidth (ANSI-stripping) agrees on wide runes. + if w := visibleWidth("\033[2m中\033[0m"); w != 2 { + t.Fatalf("visibleWidth wide = %d, want 2", w) + } +} + +func TestEditLoopCursorAccountsForWideRunes(t *testing.T) { + // Typing a CJK char must move the cursor two cells, not one, so the final + // reposition after "中" (prompt width 2 + 2 cells) lands at column 4. + var out bytes.Buffer + e := editorWithOutput("中", &out) + // Drive one render by feeding the rune then EOF (no submit needed). + if _, err := e.editLoop("> "); err == nil { + t.Fatalf("expected EOF error from truncated input") + } + s := out.String() + if !strings.Contains(s, "\033[4C") { + t.Fatalf("cursor not repositioned to column 4 for wide rune:\n%q", s) + } + if strings.Contains(s, "\033[3C") { + t.Fatalf("cursor used rune-count column 3 (wide rune miscounted):\n%q", s) + } +} + +func TestEditLoopRendersContinuationAndClears(t *testing.T) { + // "foo" + Shift+Enter + "bar" builds a two-line buffer; the continuation + // line is indented with the dim marker and each redraw clears the block. + var out bytes.Buffer + e := editorWithOutput("foo\x1b[13;2ubar\r", &out) + got, err := e.editLoop("> ") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "foo\nbar" { + t.Fatalf("submitted %q, want %q", got, "foo\nbar") + } + s := out.String() + if !strings.Contains(s, "\033[2m·\033[0m bar") { + t.Fatalf("output missing dim continuation prefix before %q:\n%q", "bar", s) + } + if !strings.Contains(s, "\033[J") { + t.Fatalf("output never clears the block with \\033[J:\n%q", s) + } +} + +// editorWithInput builds an editor whose editLoop reads the given byte stream, +// so raw-mode key handling can be exercised without a real terminal. +func editorWithInput(input string, history ...string) *replLineEditor { + reg := runtime.NewSlashRegistry() + reg.AddBuiltin(runtime.SlashCommand{Name: "model", Action: func(string) string { return "" }}) + reg.AddBuiltin(runtime.SlashCommand{Name: "models", Action: func(string) string { return "" }}) + r := strings.NewReader(input) + return newREPLLineEditor(r, bufio.NewReader(r), io.Discard, reg, history) +} + +func TestMLBufferNewlineSplitsAtCursor(t *testing.T) { + b := newMLBuffer() + b.insert("abcdef") + b.col = 3 // cursor between "abc" and "def" + b.newline() + if got := b.String(); got != "abc\ndef" { + t.Fatalf("newline split = %q, want %q", got, "abc\ndef") + } + if b.row != 1 || b.col != 0 { + t.Fatalf("cursor after newline = (%d,%d), want (1,0)", b.row, b.col) + } +} + +func TestEditLoopPlainEnterSubmits(t *testing.T) { + e := editorWithInput("abc\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "abc" { + t.Fatalf("submitted %q, want %q", got, "abc") + } +} + +func TestEditLoopShiftEnterInsertsNewline(t *testing.T) { + // CSI-u Shift+Enter is \x1b[13;2u; a bare \r then submits. + e := editorWithInput("abc\x1b[13;2udef\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "abc\ndef" { + t.Fatalf("submitted %q, want %q", got, "abc\ndef") + } +} + +func TestEditLoopModifyOtherKeysEnterInsertsNewline(t *testing.T) { + // modifyOtherKeys Shift+Enter is \x1b[27;2;13~. + e := editorWithInput("x\x1b[27;2;13~y\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "x\ny" { + t.Fatalf("submitted %q, want %q", got, "x\ny") + } +} + +func TestEditLoopCSIuPlainEnterSubmits(t *testing.T) { + // Unmodified Enter reported as CSI-u \x1b[13u must submit, not newline. + e := editorWithInput("hi\x1b[13u") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "hi" { + t.Fatalf("submitted %q, want %q", got, "hi") + } +} + +func TestEditLoopCSIuCtrlDEmptyReturnsEOF(t *testing.T) { + // Under the kitty keyboard protocol Ctrl+D on an empty line arrives as the + // CSI-u report \x1b[100;5u (code 'd', ctrl modifier) rather than raw 0x04, + // and must still exit with io.EOF. + e := editorWithInput("\x1b[100;5u") + got, err := e.editLoop("") + if !errors.Is(err, io.EOF) { + t.Fatalf("want io.EOF, got err=%v got=%q", err, got) + } +} + +func TestEditLoopCSIuCtrlDNonEmptyIgnored(t *testing.T) { + // Ctrl+D on a non-empty line is a no-op (matches the raw-byte behavior); the + // following Enter submits the typed text. + e := editorWithInput("ab\x1b[100;5u\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "ab" { + t.Fatalf("submitted %q, want %q", got, "ab") + } +} + +func TestEditLoopCSIuCtrlCInterrupts(t *testing.T) { + // Ctrl+C reported as CSI-u \x1b[99;5u must interrupt the line. + e := editorWithInput("\x1b[99;5u") + _, err := e.editLoop("") + if !errors.Is(err, errLineInterrupted) { + t.Fatalf("want errLineInterrupted, got %v", err) + } +} + +func TestMLBufferLeftRightCrossLines(t *testing.T) { + b := newMLBuffer() + b.setString("ab\ncd") // cursor at end of "cd" → (1,2) + b.left() // (1,1) + b.left() // (1,0) + if b.row != 1 || b.col != 0 { + t.Fatalf("after two lefts = (%d,%d), want (1,0)", b.row, b.col) + } + b.left() // cross to end of "ab" → (0,2) + if b.row != 0 || b.col != 2 { + t.Fatalf("left at line start = (%d,%d), want (0,2)", b.row, b.col) + } + b.right() // cross to start of "cd" → (1,0) + if b.row != 1 || b.col != 0 { + t.Fatalf("right at line end = (%d,%d), want (1,0)", b.row, b.col) + } + // Left at the very origin is a no-op. + b.setString("x") + b.home() + b.left() + if b.row != 0 || b.col != 0 { + t.Fatalf("left at origin moved cursor: (%d,%d)", b.row, b.col) + } + // Right at the very end is a no-op. + b.end() + b.right() + if b.row != 0 || b.col != 1 { + t.Fatalf("right at end moved cursor: (%d,%d)", b.row, b.col) + } +} + +func TestMLBufferUpDownClampColumn(t *testing.T) { + b := newMLBuffer() + b.setString("long line\nhi") // cursor at end of "hi" → (1,2) + b.up() // move to "long line", col stays 2 + if b.row != 0 || b.col != 2 { + t.Fatalf("up = (%d,%d), want (0,2)", b.row, b.col) + } + b.end() // col = 9 + b.down() + // Down to "hi" (len 2) clamps col from 9 to 2. + if b.row != 1 || b.col != 2 { + t.Fatalf("down clamp = (%d,%d), want (1,2)", b.row, b.col) + } + // Up on the first line is a no-op; down on the last line is a no-op. + b.setString("a\nb") + b.up() + b.up() + if b.row != 0 { + t.Fatalf("up past first line: row=%d", b.row) + } + b.down() + b.down() + if b.row != 1 { + t.Fatalf("down past last line: row=%d", b.row) + } +} + +func TestMLBufferHomeEnd(t *testing.T) { + b := newMLBuffer() + b.setString("héllo") // multi-byte, 5 runes + b.home() + if b.col != 0 { + t.Fatalf("home col = %d, want 0", b.col) + } + b.end() + if b.col != 5 { + t.Fatalf("end col = %d, want 5", b.col) + } +} + +func TestEditLoopLeftArrowThenInsert(t *testing.T) { + // Type "abc", move left once (\x1b[D), insert "X", submit → "abXc". + e := editorWithInput("abc\x1b[DX\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "abXc" { + t.Fatalf("submitted %q, want %q", got, "abXc") + } +} + +func TestEditLoopUpArrowEditsPreviousLine(t *testing.T) { + // "a" + Shift+Enter + "b", then up-arrow to line 0, End, insert "Z": + // line 0 becomes "aZ", line 1 stays "b" → "aZ\nb". + e := editorWithInput("a\x1b[13;2ub\x1b[A\x1b[FZ\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "aZ\nb" { + t.Fatalf("submitted %q, want %q", got, "aZ\nb") + } +} + +func TestMLBufferInsertMidLine(t *testing.T) { + b := newMLBuffer() + b.setString("abc") + b.home() + b.right() // cursor between "a" and "bc" + b.insert("X") + if got := b.String(); got != "aXbc" { + t.Fatalf("mid-line insert = %q, want %q", got, "aXbc") + } + if b.col != 2 { + t.Fatalf("col after insert = %d, want 2", b.col) + } +} + +func TestMLBufferBackspaceMergesLines(t *testing.T) { + b := newMLBuffer() + b.setString("ab\ncd") // cursor at (1,2) + b.home() // cursor at (1,0) + b.backspace() // merge line 1 into line 0 + if got := b.String(); got != "abcd" { + t.Fatalf("merge = %q, want %q", got, "abcd") + } + if b.row != 0 || b.col != 2 { + t.Fatalf("cursor after merge = (%d,%d), want (0,2)", b.row, b.col) + } + // Merge with a multi-byte previous line lands the cursor by rune count. + b.setString("café\nx") + b.home() + b.backspace() + if got := b.String(); got != "caféx" { + t.Fatalf("multi-byte merge = %q, want %q", got, "caféx") + } + if b.col != 4 { + t.Fatalf("cursor col after multi-byte merge = %d, want 4", b.col) + } +} + +func TestMLBufferBackspaceFirstLineCol0IsNoop(t *testing.T) { + b := newMLBuffer() + b.setString("abc") + b.home() + b.backspace() + if got := b.String(); got != "abc" { + t.Fatalf("col-0 backspace on first line = %q, want unchanged", got) + } + if b.row != 0 || b.col != 0 { + t.Fatalf("cursor moved: (%d,%d)", b.row, b.col) + } +} + +func TestEditLoopCrossLineBackspaceMerges(t *testing.T) { + // "ab" + Shift+Enter + "cd" → two lines; Home to line-1 start, backspace + // merges into "abcd", submit. + e := editorWithInput("ab\x1b[13;2ucd\x1b[H\x7f\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "abcd" { + t.Fatalf("submitted %q, want %q", got, "abcd") + } +} + +func TestEditLoopBackslashContinues(t *testing.T) { + // A line ending with a single unescaped "\" + Enter continues; a second + // line + Enter submits the two-line block without the continuation "\". + e := editorWithInput("foo\\\rbar\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "foo\nbar" { + t.Fatalf("submitted %q, want %q", got, "foo\nbar") + } +} + +func TestEditLoopEscapedBackslashSubmits(t *testing.T) { + // A line ending with "\\" (escaped) + Enter submits, keeping one literal + // backslash — it does not continue. + e := editorWithInput("foo\\\\\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "foo\\" { + t.Fatalf("submitted %q, want %q", got, "foo\\") + } +} + +func TestEditLoopBackslashMultipleContinuations(t *testing.T) { + // Three continued lines accumulate into a three-line block. + e := editorWithInput("a\\\rb\\\rc\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "a\nb\nc" { + t.Fatalf("submitted %q, want %q", got, "a\nb\nc") + } +} + +func TestMLBufferEnterContinuesTrailingRuns(t *testing.T) { + // Odd run continues and collapses to k/2 backslashes; even run submits and + // halves the run. + b := newMLBuffer() + b.setString("x\\") // one trailing backslash + if !b.enterContinues() { + t.Fatalf("single trailing backslash should continue") + } + if got := b.line(); got != "x" { + t.Fatalf("line after continue = %q, want %q", got, "x") + } + b.setString("y\\\\") // two trailing backslashes + if b.enterContinues() { + t.Fatalf("escaped double backslash should not continue") + } + if got := b.line(); got != "y\\" { + t.Fatalf("line after submit = %q, want %q", got, "y\\") + } + b.setString("z\\\\\\") // three trailing backslashes → continue, keep one + if !b.enterContinues() { + t.Fatalf("triple trailing backslash should continue") + } + if got := b.line(); got != "z\\" { + t.Fatalf("line after triple continue = %q, want %q", got, "z\\") + } +} + +func TestRememberPreservesInternalNewlines(t *testing.T) { + // A submitted multi-line block is stored as ONE history record; remember + // trims only the outer whitespace, never the internal newlines. + e := testLineEditor() + e.remember(" foo\nbar ") + if len(e.history) != 1 { + t.Fatalf("history = %v, want a single record", e.history) + } + if e.history[0] != "foo\nbar" { + t.Fatalf("remembered %q, want %q (internal newline kept)", e.history[0], "foo\nbar") + } +} + +func TestEditLoopHistoryRestoresMultiLineAndSubmits(t *testing.T) { + // Up-arrow on a blank line restores the newest history entry; a multi-line + // record comes back as a multi-line buffer that submits intact with its + // internal newline (i.e. the block is one message, not two). + e := editorWithInput("\x1b[A\r", "foo\nbar") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "foo\nbar" { + t.Fatalf("restored+submitted %q, want %q", got, "foo\nbar") + } +} + +func TestEditLoopMultiLineSubmitJoinsWithNewline(t *testing.T) { + // Three Shift+Enter lines submit as one \n-joined message, so the full + // multi-line string reaches the agent in a single turn. + e := editorWithInput("a\x1b[13;2ub\x1b[13;2uc\r") + got, err := e.editLoop("") + if err != nil { + t.Fatalf("editLoop error: %v", err) + } + if got != "a\nb\nc" { + t.Fatalf("submitted %q, want %q", got, "a\nb\nc") + } +} diff --git a/pigo/internal/cli/repl/plugin_commands_test.go b/pigo/internal/cli/repl/plugin_commands_test.go new file mode 100644 index 0000000..542e944 --- /dev/null +++ b/pigo/internal/cli/repl/plugin_commands_test.go @@ -0,0 +1,254 @@ +package repl + +// Tests for plugin slash-command wiring (#265): a plugin-declared command +// (Manager.Commands()) is registered into the REPL slash registry as a hybrid +// (Run) command, and invoking it calls Plugin.CallCommand, surfaces the +// returned notifications, and runs the returned prompt as the next turn. A real +// plugin subprocess is compiled and Discover-loaded so the JSON-RPC transport, +// handshake and commands/call round-trip are exercised end to end. + +import ( + "bytes" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + rt "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" +) + +// cmdPluginMain is a standalone plugin that declares one "hello" slash command +// and answers commands/call by echoing back a prompt built from the command's +// args plus one notification. It lets the test assert registration, notification +// surfacing, and prompt injection. It also records the raw arguments it received +// so the test can confirm a bare invocation sends a JSON string ("") not null. +const cmdPluginMain = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` + Params json.RawMessage ` + "`json:\"params\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + w := bufio.NewWriter(os.Stdout) + for sc.Scan() { + var r req + if err := json.Unmarshal(sc.Bytes(), &r); err != nil { + continue + } + switch r.Method { + case "initialize": + reply(w, r.ID, json.RawMessage(` + "`" + `{"name":"greeter","commands":[{"name":"hello","description":"say hello"}]}` + "`" + `)) + case "commands/call": + var p struct { + Name string ` + "`json:\"name\"`" + ` + Args json.RawMessage ` + "`json:\"arguments\"`" + ` + } + json.Unmarshal(r.Params, &p) + // args is a JSON string (never null); decode it to prove the contract. + var argText string + json.Unmarshal(p.Args, &argText) + res, _ := json.Marshal(map[string]any{ + "prompt": "please greet " + argText, + "notifications": []map[string]any{ + {"message": "invoked hello", "type": "info"}, + }, + }) + reply(w, r.ID, res) + case "shutdown": + return + } + } +} + +func reply(w *bufio.Writer, id *json.RawMessage, result json.RawMessage) { + if id == nil { + return + } + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() +} +` + +// buildPluginInDir compiles src into an executable named bin inside dir and +// returns nothing (the executable path is dir/bin). Discover loads any +// executable regular file directly under dir. +func buildPluginInDir(t *testing.T, dir, bin, src string) { + t.Helper() + srcPath := filepath.Join(dir, "plugin_main.go") + if err := os.WriteFile(srcPath, []byte(src), 0o644); err != nil { + t.Fatalf("write plugin source: %v", err) + } + binPath := filepath.Join(dir, bin) + if runtime.GOOS == "windows" { + binPath += ".exe" + } + cmd := exec.Command("go", "build", "-o", binPath, srcPath) + cmd.Env = os.Environ() + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("build plugin: %v\n%s", err, out) + } + // Remove the source so Discover only sees the executable (a .go file is not + // executable, but keeping the dir clean avoids any ambiguity). + _ = os.Remove(srcPath) +} + +// loadTestManager compiles the greeter plugin into a fresh dir and Discover-loads +// it, returning a Manager with exactly that one plugin. The caller must Close it. +func loadTestManager(t *testing.T) *plugin.Manager { + t.Helper() + // Build in a build dir, then move only the binary into the plugins dir so + // Discover (which loads every executable file in the dir) sees just the one + // plugin executable. + buildDir := t.TempDir() + buildPluginInDir(t, buildDir, "greeter", cmdPluginMain) + pluginsDir := t.TempDir() + binName := "greeter" + if runtime.GOOS == "windows" { + binName += ".exe" + } + if err := os.Rename(filepath.Join(buildDir, binName), filepath.Join(pluginsDir, binName)); err != nil { + t.Fatalf("move plugin binary: %v", err) + } + mgr, err := plugin.Discover(pluginsDir, os.Stderr, os.Stderr) + if err != nil { + t.Fatalf("Discover: %v", err) + } + if len(mgr.Commands()) != 1 || mgr.Commands()[0].Spec.Name != "hello" { + t.Fatalf("expected one discovered command 'hello', got %+v", mgr.Commands()) + } + return mgr +} + +// TestBuildSlashRegistryRegistersPluginCommand verifies a discovered plugin +// command is registered as a resolvable slash command in the registry. +func TestBuildSlashRegistryRegistersPluginCommand(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + mgr := loadTestManager(t) + defer mgr.Close() + + reg, err := prompts.BuildSlashRegistry(&cli.LiveConfig{Model: "faux", ProviderName: "faux"}, nil, mgr, prompts.PromptTemplateSources{}) + if err != nil { + t.Fatalf("buildSlashRegistry: %v", err) + } + if _, ok := reg.Lookup("hello"); !ok { + t.Fatalf("plugin command /hello was not registered in the slash registry") + } + + // Resolving it must run the plugin (side effect), surface its notification, + // and yield the plugin's prompt to run. + out, err := reg.ResolveOutcome("/hello world") + if err != nil { + t.Fatalf("ResolveOutcome(/hello): %v", err) + } + if !out.Handled || out.Kind != rt.SlashPrompt { + t.Fatalf("outcome = %+v, want handled SlashPrompt", out) + } + if !strings.Contains(out.Message, "invoked hello") { + t.Errorf("notification not surfaced in Message: %q", out.Message) + } + if out.Prompt != "please greet world" { + t.Errorf("Prompt = %q, want plugin-returned prompt", out.Prompt) + } +} + +// TestBuiltinWinsOverPluginCommand verifies a built-in command of the same name +// wins over a plugin command (existing precedence preserved): the plugin command +// is shadowed and the built-in's behavior is what resolves. +func TestBuiltinWinsOverPluginCommand(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + mgr := loadTestManager(t) + defer mgr.Close() + + reg := rt.NewSlashRegistry() + reg.AddBuiltin(rt.SlashCommand{ + Name: "hello", + Action: func(string) string { return "builtin hello" }, + }) + prompts.RegisterPluginCommands(reg, mgr) + + if names := reg.Shadowed(); len(names) != 1 || names[0].Name != "hello" { + t.Fatalf("plugin command should be shadowed by built-in, shadowed=%v", names) + } + out, err := reg.ResolveOutcome("/hello there") + if err != nil { + t.Fatalf("ResolveOutcome: %v", err) + } + if out.Kind != rt.SlashAction || out.Message != "builtin hello" { + t.Errorf("built-in must win: outcome = %+v", out) + } +} + +// TestREPLPluginCommandInjectsPrompt drives the full REPL: invoking a plugin +// slash command prints its notification and runs the returned prompt as the next +// agent turn (so the fake provider is called once and the injected prompt lands +// in the conversation history). +func TestREPLPluginCommandInjectsPrompt(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + mgr := loadTestManager(t) + defer mgr.Close() + + reg, err := prompts.BuildSlashRegistry(&cli.LiveConfig{Model: "faux", ProviderName: "faux"}, nil, mgr, prompts.PromptTemplateSources{}) + if err != nil { + t.Fatalf("buildSlashRegistry: %v", err) + } + + store, err := session.NewStore(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + p := &replProvider{reply: "hi there"} + live := &cli.LiveConfig{Model: "faux", ProviderName: "faux", Provider: p} + deps := replDeps{ + store: store, + header: session.SessionHeader{ID: session.NewID(time.Now().UTC()), Model: "faux", Provider: "faux"}, + agentCtx: &agentcore.AgentContext{}, + live: live, + reg: agenttool.NewToolRegistry(), + slash: reg, + creds: provider.NewCredentialStore(nil), + } + + var out bytes.Buffer + if err := runREPL(strings.NewReader("/hello world\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + + // The plugin's notification must have been printed. + if !strings.Contains(out.String(), "invoked hello") { + t.Errorf("plugin notification not printed, out=%q", out.String()) + } + // The returned prompt must have run exactly one turn. + if p.calls != 1 { + t.Fatalf("plugin command should inject and run exactly 1 turn, got %d", p.calls) + } + // The injected prompt must be the user message that started the turn. + if len(deps.agentCtx.Messages) == 0 { + t.Fatal("expected messages in context after the injected turn") + } + u0, ok := deps.agentCtx.Messages[0].(agentcore.UserMessage) + if !ok || agentcore.ContentToText(u0.Content) != "please greet world" { + t.Errorf("injected prompt not run as the turn: %+v", deps.agentCtx.Messages[0]) + } +} diff --git a/pigo/internal/cli/repl/remotecontrol.go b/pigo/internal/cli/repl/remotecontrol.go new file mode 100644 index 0000000..dbf6f07 --- /dev/null +++ b/pigo/internal/cli/repl/remotecontrol.go @@ -0,0 +1,237 @@ +// This file wires the remote-control bridge (internal/remotecontrol, #442) into +// the interactive REPL (#443). It adds the "/remote-control" command that +// starts/stops an in-process HTTP+WebSocket server mirroring the session to a +// paired browser on the LAN, tees REPL output to that browser, merges +// browser-submitted prompts into the input loop, and routes tool-call +// confirmations to the browser while a client is connected. +// +// The design keeps the non-remote path byte-identical: when no remote session +// is active, teeWriter forwards only to stdout, the input select degenerates to +// a plain editor read (the remote channel is nil), and confirmations use the +// local stdin prompt unchanged. +package repl + +import ( + "context" + "fmt" + "io" + "strings" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/remotecontrol" + "github.com/smallnest/pigo/internal/trust" +) + +// teeWriter is an io.Writer that always forwards to a primary writer (the +// terminal) and, when a secondary is set, mirrors the same bytes to it (the +// remote browser). It is safe for concurrent Write/setSecondary because the +// bridge output writer is fed from the REPL goroutine while the secondary is +// toggled by the /remote-control command on the same goroutine, but writes to +// the WebSocket happen on other goroutines; the mutex keeps the swap atomic. +type teeWriter struct { + primary io.Writer + mu sync.Mutex + second io.Writer +} + +func newTeeWriter(primary io.Writer) *teeWriter { return &teeWriter{primary: primary} } + +func (t *teeWriter) setSecondary(w io.Writer) { + t.mu.Lock() + t.second = w + t.mu.Unlock() +} + +func (t *teeWriter) Write(p []byte) (int, error) { + // The primary write is authoritative for the returned count/err so terminal + // behavior is unchanged; a mirror failure never breaks the local session. + n, err := t.primary.Write(p) + t.mu.Lock() + second := t.second + t.mu.Unlock() + if second != nil { + _, _ = second.Write(p) + } + return n, err +} + +// remoteSession owns the running server + bridge for one /remote-control +// activation. It is nil in deps until the command starts a session, and is +// cleared on stop. +type remoteSession struct { + server *remotecontrol.Server + bridge *remotecontrol.Bridge + url string +} + +// inputChan returns the bridge's remote-input channel while a session is +// active, or nil when inactive. A nil channel blocks forever in a select, so +// the input loop transparently ignores remote input when remote control is off. +func (rs *remoteSession) inputChan() <-chan string { + if rs == nil || rs.bridge == nil { + return nil + } + return rs.bridge.RemoteInput() +} + +// hasClient reports whether a browser is currently paired and connected. +func (rs *remoteSession) hasClient() bool { + return rs != nil && rs.bridge != nil && rs.bridge.Enabled() +} + +// runRemoteControl handles the "/remote-control" command and its "stop"/"status" +// subcommands. It mutates deps in place (deps.remote, deps.tee) so the input +// loop and output tee pick up the change on the next iteration. +func runRemoteControl(out io.Writer, deps *replDeps, line string) { + arg := strings.TrimSpace(strings.TrimPrefix(line, "/remote-control")) + switch arg { + case "stop": + stopRemoteControl(out, deps) + case "status", "": + if arg == "status" { + remoteControlStatus(out, deps) + return + } + startRemoteControl(out, deps) + default: + fmt.Fprintf(out, "usage: /remote-control [stop|status]\n") + } +} + +func startRemoteControl(out io.Writer, deps *replDeps) { + if deps.remote != nil { + fmt.Fprintf(out, "remote control already running: %s\n", deps.remote.url) + return + } + // Handler is set after the server is built (SetHandler), but NewServer takes + // it up front; the bridge's Sink is the server itself, so build the server + // first with the bridge as handler once the bridge exists. To break the + // cycle we construct the server, then the bridge (Sink=server), then tell the + // server to route client frames to the bridge. + // The connect/disconnect callbacks print a terminal notice so the operator + // sees when a browser gains or loses remote access to this session (§7.3). + // They run on the server's WebSocket goroutine and only write a line, so they + // don't block. + cfg := remotecontrol.Config{ + OnClientConnect: func(remoteAddr string) { + fmt.Fprintf(out, "\n[remote-control] browser connected from %s\n", remoteAddr) + }, + OnClientDisconnect: func() { + fmt.Fprintf(out, "\n[remote-control] browser disconnected\n") + }, + } + srv := remotecontrol.NewServer(cfg, nil) + bridge := remotecontrol.NewBridge(srv) + srv.SetHandler(bridge) + + url, err := srv.Start() + if err != nil { + fmt.Fprintf(out, "remote control: %v\n", err) + return + } + rs := &remoteSession{server: srv, bridge: bridge, url: url} + deps.remote = rs + if deps.tee != nil { + deps.tee.setSecondary(bridge.OutputWriter()) + } + + fmt.Fprintf(out, "\nRemote control started. Open this URL on a device on the same network:\n\n %s\n\n", url) + if qr, qerr := remotecontrol.Render(url); qerr == nil { + fmt.Fprintln(out, qr) + } + fmt.Fprintln(out, "Run /remote-control stop to end the session.") +} + +func stopRemoteControl(out io.Writer, deps *replDeps) { + if deps.remote == nil { + fmt.Fprintln(out, "remote control is not running") + return + } + if deps.tee != nil { + deps.tee.setSecondary(nil) + } + _ = deps.remote.server.Stop(context.Background()) + deps.remote = nil + fmt.Fprintln(out, "remote control stopped") +} + +func remoteControlStatus(out io.Writer, deps *replDeps) { + if deps.remote == nil { + fmt.Fprintln(out, "remote control: off") + return + } + state := "waiting for a browser to connect" + if deps.remote.hasClient() { + state = "browser connected" + } + fmt.Fprintf(out, "remote control: on (%s)\n %s\n", state, deps.remote.url) +} + +// beforeToolCall builds the tool-call confirmation seam for a turn. It always +// constructs the local stdin prompt (trust.BeforeToolCall) and, when a remote +// session exists, wraps it with bridgeBeforeToolCall so confirmations route to a +// paired browser while one is connected. When deps.remote is nil the wrapper is +// skipped entirely, so the returned func is exactly the local seam — the +// non-remote path is byte-identical to before (#443). +func beforeToolCall(deps replDeps, out io.Writer) agentcore.BeforeToolCallFunc { + local := trust.BeforeToolCall(deps.trust, deps.cwd, deps.in, out, deps.confirmMu) + if deps.remote == nil { + return local + } + return bridgeBeforeToolCall(deps.trust, deps.cwd, deps.remote, out, deps.confirmMu, local) +} + +// bridgeBeforeToolCall wraps the local stdin confirmation seam so that while a +// browser is connected, side-effect tool-call confirmations are routed to the +// browser instead of blocking on the local terminal. When no browser is +// connected it delegates to the local prompt so behavior is unchanged. +// +// This mirrors trust.BeforeToolCall's gating (side-effect tools only, honoring +// session trust) but delegates the allow/always decision to the remote client +// via Bridge.Confirm. A ctx cancellation (e.g. SIGINT) makes Confirm return +// remote=false, which we treat as a denial so an interrupted run does not +// silently proceed. Refinements (local-answer race, timeouts) are the hardening +// node's job (#445). +func bridgeBeforeToolCall(mgr *trust.Manager, cwd string, rs *remoteSession, out io.Writer, mu *sync.Mutex, local agentcore.BeforeToolCallFunc) agentcore.BeforeToolCallFunc { + return func(ctx context.Context, call agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { + if !rs.hasClient() || mgr == nil { + if local != nil { + return local(ctx, call) + } + return nil + } + if !trust.SideEffectTools[call.Name] { + return nil + } + if mu != nil { + mu.Lock() + defer mu.Unlock() + } + if mgr.IsTrusted(cwd) { + return nil + } + summary := trust.ToolCallSummary(call) + fmt.Fprintf(out, "\npigo wants to run %q — approve on the paired device…\n", call.Name) + d, remote := rs.bridge.Confirm(ctx, call.Name, summary) + if !remote { + // Interrupted / cancelled before the browser answered: deny. + return blockToolCall(call, cwd) + } + if d.Always { + mgr.SetSessionTrust(cwd) + } + if !d.Approve { + return blockToolCall(call, cwd) + } + return nil + } +} + +func blockToolCall(call agentcore.AgentToolCall, cwd string) *agentcore.BeforeToolCallDecision { + msg := fmt.Sprintf("tool %q blocked: %s is not trusted (use /trust to trust this project)", call.Name, cwd) + return &agentcore.BeforeToolCallDecision{ + Block: true, + Content: &agentcore.ContentList{agentcore.NewTextContent(msg)}, + } +} diff --git a/pigo/internal/cli/repl/repl.go b/pigo/internal/cli/repl/repl.go new file mode 100644 index 0000000..08f0802 --- /dev/null +++ b/pigo/internal/cli/repl/repl.go @@ -0,0 +1,1150 @@ +// This file implements the line-based interactive REPL (US-003/#106) that +// replaces the former full-screen bubbletea TUI. When pigo is invoked without a +// prompt on a terminal it runs a simple read → run → stream-print loop over a +// persisted session: no full-screen rendering, no popup menus, no viewport — a +// prompt, the agent's streamed reply, and a new prompt. +// +// The REPL is a synchronous loop on the main goroutine: it reads one line, +// resolves slash-commands, launches an agent run, drains the run's event stream +// to stdout as it streams, persists the session, and returns to the prompt. A +// SIGINT during a run cancels that run's context and returns to the prompt; when +// idle, the reader sees EOF/interrupt and exits cleanly. +package repl + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "os" + "os/signal" + "strconv" + "strings" + "sync" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/btw" + "github.com/smallnest/pigo/internal/cli/goal" + "github.com/smallnest/pigo/internal/cli/memstatus" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/status" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/clipboard" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/memory" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +// replDeps bundles the collaborators a REPL run needs. They are assembled once +// by Run and reused across every prompt in the session. +type replDeps struct { + store *session.Store + header session.SessionHeader + agentCtx *agentcore.AgentContext + live *cli.LiveConfig + reg *agenttool.ToolRegistry + // reminders holds the per-turn system-reminder providers (US-002). It is nil + // when no todo tool is present; when set, streamRun wires it so ephemeral + // context is injected into each turn's request. + reminders *runtime.ReminderRegistry + slash *runtime.SlashRegistry + creds *provider.CredentialStore + + // notifier delivers agent lifecycle events to subscribed plugins (US-017, + // #133). It is nil when no plugin subscribes; DrainStream's OnEvent stays + // unset in that case. + notifier *plugin.EventNotifier + + // trust persists project-trust decisions (US-018, #134). It is nil when + // trust is disabled (e.g. the store could not be loaded); when nil the + // BeforeToolCall hook is not installed and the first-run prompt is skipped. + trust *trust.Manager + // cwd is the directory pigo was launched in, used as the trust key and as + // the directory side-effect tools are gated against. It does not change + // during a session (pigo does not cd). + cwd string + // in is the shared buffered input reader. The main loop and the tool-call + // confirmation prompt both read from it so input typed ahead is never split + // between them. It is created by Run (wrapping os.Stdin) or, for + // direct test callers, lazily from the in argument at the top of runREPL. + in *bufio.Reader + // confirmMu serializes tool-call confirmation prompts so concurrent + // parallel side-effect tool calls do not interleave on the shared + // stdin/stdout. It is a pointer so every value copy of replDeps (passed to + // streamRun per prompt) shares one mutex. + confirmMu *sync.Mutex + + // curLeaf is the id of the entry the conversation currently descends from — + // the active leaf of the on-disk session tree (US-007, #123). A fresh session + // starts empty (""); each persisted turn advances it to the newly written + // leaf. /tree can move it to a historical entry so the next turn branches from + // there rather than the tip. + curLeaf string + // persisted is the number of agentCtx.Messages already written to disk. The + // per-turn persist appends only Messages[persisted:] as a branch descending + // from curLeaf, so switching leaves and continuing grows a real tree instead + // of rewriting the file as a single linear chain. + persisted int + + // goal holds the session's autonomous-goal state (mirrors pi-goal), driven by + // the /goal command. It is always non-nil (Run seeds an idle + // state); the GoalReminderProvider and the goal tools share this handle so + // the objective is injected each turn and goal_complete/goal_blocked can end + // the run. + goal *agenttool.GoalState + + // lastBtw holds the most recent /btw side thread from this process (US-004, + // #281), so a bare "/btw" reopens it and shows its history. It is nil until + // the first /btw and is reset to nil on any session switch (/fork, /clone, + // /import) because a side thread branched from the old conversation no longer + // makes sense against a different one. It is never persisted: restarting pigo + // starts with lastBtw nil again. + lastBtw *agentcore.AgentContext + // lastBtwBase is the number of leading messages in lastBtw.Messages that were + // copied from the main conversation as background — the side thread's own Q&A + // is everything from this index on. Reopening replays only Messages[base:] so + // the whole main transcript is not reprinted. + lastBtwBase int + + // telemetry holds the retained per-run telemetry events (US-001, #291) and + // the cumulative accumulator that sums metrics across all runs in the session. + // It is reset to "no telemetry yet" on any session switch (/fork, /clone, + // /import). + telemetry *cli.TelemetryHolder + + // dispatcher is the run's hook dispatcher (#425), nil when no hooks are + // configured (FR-18) so every hook path is skipped. It is resolved once per + // session (trust-gated, FR-14) at the top of runREPL, which also fires + // SessionStart. Each turn's streamRun installs the per-turn seams + // (PreToolUse/PostToolUse/Stop) and runs UserPromptSubmit through it. + dispatcher *hooks.Dispatcher + // hookDeps carries the session id / project dir stamped onto every HookInput. + hookDeps run.HookDeps + + // remote holds the active remote-control session (#443), or nil when remote + // control is off. The /remote-control command sets/clears it in place; the + // input loop selects on its RemoteInput channel and the confirm seam routes + // tool-call approvals to the paired browser while it is non-nil. When nil the + // REPL behaves exactly as before (byte-identical). + remote *remoteSession + // tee wraps out so REPL output can be mirrored to the remote browser while a + // session is active. It is created at the top of runREPL; the /remote-control + // command toggles its secondary writer. When no remote session is active it + // forwards only to the terminal. + tee *teeWriter + + // memoryRoot is the persistent-memory Store root (empty when memory is + // disabled). streamRun routes auto-compaction checkpoints here and /rebuild + // recovers from /sessions//, the canonical checkpoint location. + memoryRoot string + // memstore is the live persistent-memory Store (nil when memory is disabled). + // It lets /memory inspect entry counts without re-opening the database. + memstore *memory.Store + + // snap is the shared file-snapshot recorder backing /rewind (nil under + // --no-tools). Each turn's write/edit mutations are captured into it; after the + // turn persists, commitRewindPoint groups them into a restore point tagged with + // the pre-turn leaf so /rewind can roll files and the conversation back together. + snap *agenttool.FileSnapshotRecorder + + // jobs holds background bash jobs launched with run_in_background. On REPL + // exit its still-running jobs are killed so background processes are not + // orphaned. nil when the shell tool is disabled. + jobs *agenttool.BashJobStore +} + +// replScanBufInit is the initial size of the shared input reader. A REPL user +// may paste a long single line (a big prompt or a pasted file); bufio.Reader +// grows its returned string beyond this buffer on demand, so a long line is +// still read whole (just accumulated rather than capped). This drops the hard +// 4MiB cap the previous bufio.Scanner had: ReadString is unbounded, so a +// pathological paste can grow the line buffer. That is an acceptable tradeoff +// for a terminal REPL (where OS line buffering bounds normal input) in exchange +// for correct buffer sharing with the tool-call confirmation prompt. +const replScanBufInit = 64 * 1024 + +// notifierHandle returns the plugin event-delivery callback for this session's +// runs, or nil when no plugin subscribed. Returning nil (rather than a func that +// dispatches to a nil notifier) keeps DrainStream's OnEvent unset in the common +// no-plugin case, so the drain loop skips the per-event call entirely. +func (deps replDeps) notifierHandle() func(agentcore.AgentEvent) { + if deps.notifier == nil { + return nil + } + return deps.notifier.Handle +} + +// runREPL runs the read → run → stream-print loop until EOF, /exit or /quit. It +// reads from in (os.Stdin in production) and writes prompts, streamed replies +// and status lines to out (os.Stdout). It is the interactive replacement for the +// bubbletea program. +func runREPL(in io.Reader, out io.Writer, deps replDeps) error { + // Use the shared buffered reader so the main loop and the tool-call + // confirmation prompt read from the same buffer: input is never split + // between them the way it would be if the loop used a bufio.Scanner with + // its own private buffer (a Scanner can read past the current line into its + // internal buffer, trapping bytes the confirmation prompt would then never + // see). deps.in is set by Run (wrapping os.Stdin); the in + // parameter is only a fallback for direct callers (tests) that do not set + // deps.in, and is ignored once deps.in is populated. + if deps.in == nil { + deps.in = bufio.NewReaderSize(in, replScanBufInit) + } + + // Mirror all REPL output to the remote browser while a /remote-control + // session is active (#443). teeWriter forwards to the terminal + // unconditionally and, when the command toggles a secondary, also to the + // bridge's output writer. When remote control is off it is a thin + // pass-through, so terminal output is byte-identical to before. + deps.tee = newTeeWriter(out) + out = deps.tee + + // Stop a still-running remote-control server when the REPL exits (FR-16), so + // quitting pigo tears down the LAN listener rather than leaking it. The + // closure reads deps.remote at exit time, so it covers a session started mid + // run; an explicit /remote-control stop clears deps.remote and makes this a + // no-op. + defer func() { + if deps.remote != nil { + _ = deps.remote.server.Stop(context.Background()) + deps.remote = nil + } + }() + + // Kill any still-running background bash jobs on exit so long-running + // commands (dev servers, watchers) launched with run_in_background are not + // orphaned when the REPL quits. No-op when the shell tool is disabled. + defer func() { + if deps.jobs != nil { + deps.jobs.KillAll() + } + }() + + // Resolve the run's hooks once per session and fire SessionStart so any + // injected context lands in the first turn (#423/#425). The project layer is + // honored only when the working directory is trusted (FR-14). A malformed hook + // layer disables hooks with a warning rather than aborting an interactive + // session. deps.dispatcher stays nil when no hooks are configured (FR-18), so + // streamRun skips all hook work. + deps.hookDeps = run.HookDeps{SessionID: deps.header.ID, ProjectDir: deps.cwd, WarnLog: out} + trusted := deps.trust != nil && deps.trust.IsTrusted(deps.cwd) + if set, err := run.ResolveHookSet(deps.cwd, trusted); err != nil { + fmt.Fprintf(out, "pigo: hooks disabled: %v\n", err) + } else if d := run.BuildDispatcher(set, deps.hookDeps); d != nil { + deps.dispatcher = d + if deps.reminders == nil { + deps.reminders = runtime.NewReminderRegistry() + } + ssCfg := runtime.RunConfig{Reminders: deps.reminders} + run.DispatchSessionStart(context.Background(), d, &ssCfg, deps.hookDeps, "startup") + deps.reminders = ssCfg.Reminders + } + var priorInputs []string + for _, msg := range deps.agentCtx.Messages { + if user, ok := msg.(agentcore.UserMessage); ok { + priorInputs = append(priorInputs, agentcore.ContentToText(user.Content)) + } + } + editor := newREPLLineEditor(in, deps.in, out, deps.slash, priorInputs) + editor.models = append([]string{deps.live.Model}, editor.models...) + + // A SIGINT during a run cancels only that run; the handler is installed for + // the whole REPL and targets whichever run is active via runCancel. runCancel + // is read on the signal goroutine and written on the main loop, so a mutex + // guards it against a data race. + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, os.Interrupt) + defer signal.Stop(sigCh) + var ( + mu sync.Mutex + runCancel context.CancelFunc + ) + setCancel := func(c context.CancelFunc) { + mu.Lock() + runCancel = c + mu.Unlock() + } + go func() { + for range sigCh { + mu.Lock() + c := runCancel + mu.Unlock() + if c != nil { + c() + } + } + }() + + // Input is acquired by a dedicated reader goroutine so the main loop can + // select between a locally-typed line and a line submitted from the paired + // browser (#443). The goroutine reads one line per request sent on promptReq + // and returns it on localCh; the loop only issues a new request when the + // reader is idle (readerBusy == false), so a browser-submitted line that + // arrives while a local ReadLine is still blocked does not deadlock or + // double-prompt. When remote control is off, remoteCh is nil (blocks + // forever), so the select degenerates to a plain local read — behavior is + // identical to before. + type lineResult struct { + raw string + err error + } + promptReq := make(chan string) + localCh := make(chan lineResult, 1) + go func() { + for p := range promptReq { + raw, err := editor.ReadLine(p) + localCh <- lineResult{raw: raw, err: err} + } + }() + defer close(promptReq) + readerBusy := false + + for { + replPrompt := fmt.Sprintf("pigo(%s)> ", deps.live.Model) + if !readerBusy { + fmt.Fprintln(out) + promptReq <- replPrompt + readerBusy = true + } + var ( + raw string + err error + fromRemote bool + ) + select { + case res := <-localCh: + readerBusy = false + raw, err = res.raw, res.err + case rl := <-deps.remote.inputChan(): + // A browser-submitted line. The pending local ReadLine stays blocked + // (readerBusy remains true) and will be consumed on a later iteration. + raw, fromRemote = rl, true + // Echo the remote line locally so the terminal (and, via the tee, the + // browser) shows what was submitted, mirroring a typed prompt. + fmt.Fprintf(out, "%s%s\n", replPrompt, raw) + } + if errors.Is(err, errLineInterrupted) { + continue + } + if err != nil && raw == "" { + // EOF (Ctrl+D) or read error with no partial line: exit cleanly. + fmt.Fprintln(out) + if err == io.EOF { + return nil + } + return err + } + line := strings.TrimSpace(raw) + if !fromRemote { + editor.remember(line) + } + if line == "" { + if err != nil { + // A trailing partial line at EOF that trims to empty: exit. + fmt.Fprintln(out) + return nil + } + continue + } + if line == "/exit" || line == "/quit" { + return nil + } + if line == "/compact" { + // /compact is intercepted here (like /exit) because compaction must run + // an agent stream and mutate the shared context — neither of which a + // slash Action closure (string in, string out) can do. Compaction + // replaces the whole message list with a summary + tail, so the session + // is rewritten linearly (Save) and the branch-tracking state is reset to + // the new flattened leaf. + runManualCompact(out, deps) + deps.header.UpdatedAt = time.Now().UTC() + if err := deps.store.Save(deps.header, deps.agentCtx.Messages); err != nil { + fmt.Fprintf(out, "pigo: session save failed: %v\n", err) + } + deps.persisted = len(deps.agentCtx.Messages) + deps.curLeaf = "" + if _, entries, err := deps.store.LoadEntries(deps.header.ID); err == nil && len(entries) > 0 { + deps.curLeaf = entries[len(entries)-1].ID + } + continue + } + if line == "/rebuild" { + // /rebuild is intercepted here for the same reason as /compact: it + // reconstructs the whole message list (checkpoint summary + retained + // tail) and mutates the shared context, which a slash Action closure + // cannot do. It reloads a persisted checkpoint when present, else falls + // back to lossy compaction. Like /compact, the session is rewritten + // linearly (Save) and the branch cursor reset to the new flattened leaf. + runManualRebuild(out, deps) + deps.header.UpdatedAt = time.Now().UTC() + if err := deps.store.Save(deps.header, deps.agentCtx.Messages); err != nil { + fmt.Fprintf(out, "pigo: session save failed: %v\n", err) + } + deps.persisted = len(deps.agentCtx.Messages) + deps.curLeaf = "" + if _, entries, err := deps.store.LoadEntries(deps.header.ID); err == nil && len(entries) > 0 { + deps.curLeaf = entries[len(entries)-1].ID + } + continue + } + if line == "/clone" || line == "/fork" || strings.HasPrefix(line, "/fork ") { + // /fork and /clone are intercepted here (like /compact) because they + // switch the active session — replacing the header and the shared + // context in place — which a slash Action closure (pure string→string) + // cannot do. runForkClone saves the current session, forks it, and + // swaps deps.header / deps.agentCtx to the new branch on success. + runForkClone(out, &deps, line) + continue + } + if line == "/tree" || strings.HasPrefix(line, "/tree ") { + // /tree is intercepted here (like /fork) because "/tree " moves the + // active leaf and rebuilds the shared context in place — mutating + // per-run state a slash Action closure cannot reach. With no argument + // it just prints the tree. + runTree(out, &deps, line) + continue + } + if line == "/rewind" || strings.HasPrefix(line, "/rewind ") { + // /rewind is intercepted here (like /tree) because "/rewind " restores + // files on disk AND moves the active leaf, rebuilding the shared context in + // place — per-run state a slash Action closure cannot reach. With no + // argument it lists the restore points. + runRewind(out, &deps, line) + continue + } + if line == "/export" || strings.HasPrefix(line, "/export ") { + // /export writes the current session to a file. It is intercepted here + // (not a slash Action) because it must first persist the live turn so the + // export reflects unsaved messages. The exact-or-space-prefix guard keeps + // "/exporter" from being mistaken for "/export". + runExport(out, &deps, line) + continue + } + if line == "/import" || strings.HasPrefix(line, "/import ") { + // /import loads a JSONL export as a fresh session and switches to it, + // swapping deps.header / deps.agentCtx in place — which a slash Action + // closure cannot do. The guard keeps "/important" from matching. + runImport(out, &deps, line) + continue + } + if line == "/copy" { + // /copy writes the most recent assistant text to the clipboard. It is + // intercepted here (not a slash Action) because it must read the live + // message list, which an Action closure cannot reach. + runCopy(out, &deps) + continue + } + if line == "/session" { + // /session prints live session stats (message count, tokens, compactions) + // derived from deps.header + the in-memory context — state a pure + // string→string Action closure cannot see. + runSession(out, &deps) + continue + } + if line == "/status" || strings.HasPrefix(line, "/status ") { + // /status prints a colored multi-section status report with runtime + // config, context usage, and more — state a pure string→string Action + // closure cannot see. The exact-or-space-prefix guard keeps "/statusfoo" + // from being mistaken for "/status". + status.RunStatus(out, &deps) + continue + } + if line == "/memory" || strings.HasPrefix(line, "/memory ") { + // /memory prints the persistent-memory + infinite-context report. + // Like /status it needs live state (the memory store, memory root, + // session id, and messages) that a string→string Action closure + // cannot reach, so it is intercepted here. + memstatus.RunMemory(out, deps.memstore, deps.memoryRoot, deps.header.ID, + deps.agentCtx.Messages, deps.live.ContextWindow) + continue + } + if line == "/dream" || strings.HasPrefix(line, "/dream ") { + // /dream is intercepted here (like /memory/status) because it spawns the + // process-isolated consolidation subprocess (pigo --dream) and renders the + // returned Report against live state (deps.cwd / deps.memoryRoot) — work a + // string→string Action closure cannot do. It never mutates the shared + // context, so no session re-save/leaf reset is needed. The + // exact-or-space-prefix guard keeps "/dreamer" from matching, and + // "/dream --dry-run" runs the same analysis without writing. + runDream(out, deps, line) + continue + } + if line == "/goal" || strings.HasPrefix(line, "/goal ") { + // /goal is intercepted here (like /compact) because it must run one or + // more agent streams and mutate the shared context/goal state — none of + // which a slash Action closure (string→string) can do. It drives the + // autonomous goal loop, reusing the same SIGINT cancel plumbing as a + // normal turn via setCancel. + goal.RunGoal(setCancel, out, &deps, line) + continue + } + if line == "/btw" || strings.HasPrefix(line, "/btw ") { + // /btw is intercepted here (like /goal) because it must run an agent + // stream against a COPY of the main context and must NOT mutate or + // persist the main conversation — none of which a slash Action closure + // can express. The exact-or-space-prefix guard keeps "/btweak" from + // being mistaken for "/btw". It reuses the same SIGINT cancel plumbing + // as a normal turn via setCancel. + btw.RunBtw(setCancel, out, &deps, editor, line) + continue + } + if line == "/remote-control" || strings.HasPrefix(line, "/remote-control ") { + // /remote-control is intercepted here (not a slash Action) because it + // starts/stops the in-process server and toggles deps.remote / the + // output tee in place — per-session state a pure string→string Action + // closure cannot reach (#443). The exact-or-space-prefix guard keeps a + // longer command from being mistaken for it. + runRemoteControl(out, &deps, line) + continue + } + + // Resolve slash-commands: an action command runs and prints its message + // (no agent run); a prompt command or skill expands to the text we run; a + // hybrid (plugin) command runs its side effect, prints its message + // (notifications), then runs the returned prompt if non-empty; an unknown + // command prints an error and returns to the prompt. + prompt := line + if strings.HasPrefix(line, "/") { + outcome, err := deps.slash.ResolveOutcome(line) + if err != nil { + fmt.Fprintf(out, "%v\n", err) + continue + } + if outcome.Kind == runtime.SlashAction { + if outcome.Message != "" { + fmt.Fprintln(out, outcome.Message) + } + continue + } + // SlashPrompt. A hybrid command may carry a Message to surface first + // (e.g. plugin notifications) alongside the prompt to run. + if outcome.Message != "" { + fmt.Fprintln(out, outcome.Message) + } + // A hybrid command with no prompt (only notifications) has nothing to + // run; return to the prompt without starting a turn. + if outcome.Prompt == "" { + continue + } + prompt = outcome.Prompt + } + + // Launch the run and stream it to out. runCancel is published so the SIGINT + // handler can interrupt this run; it is cleared when the run settles. + runCtx, cancel := context.WithCancel(context.Background()) + setCancel(cancel) + // Capture the leaf the turn descends from before it advances, so a rewind + // to this turn can move the conversation back to exactly here. + preTurnLeaf := deps.curLeaf + streamRun(runCtx, out, deps, prompt) + cancel() + setCancel(nil) + + // Persist the turn as a branch descending from the active leaf, so a + // prior /tree leaf-switch produces a real sibling branch on disk rather + // than a truncated linear rewrite. + deps.header.Model = deps.live.Model + deps.header.Provider = deps.live.ProviderName + cli.PersistTurn(out, &deps) + // Group this turn's file mutations (if any) into a rewind restore point. + deps.snap.Commit(preTurnLeaf, rewindLabel(prompt)) + } +} + +// streamRun appends the prompt to the shared context, starts an agent run, and +// prints the streamed assistant text and tool activity to out. It blocks until +// the run ends. The context grows in place so the next prompt continues the +// conversation. +func streamRun(ctx context.Context, out io.Writer, deps replDeps, prompt string) { + content, err := ui.BuildUserContent(prompt) + if err != nil { + fmt.Fprintf(out, "pigo: %v\n", err) + return + } + // UserPromptSubmit runs before the prompt is committed to the shared context, + // so a block aborts the turn without leaving a dangling user message; any + // additionalContext is injected into this turn only (one-shot reminder). + if deps.dispatcher != nil { + pc := runtime.RunConfig{Reminders: deps.reminders} + if block, reason := run.DispatchUserPromptSubmit(ctx, deps.dispatcher, &pc, deps.hookDeps, prompt); block { + fmt.Fprintf(out, "pigo: prompt blocked by hook: %s\n", reason) + return + } + deps.reminders = pc.Reminders + } + deps.agentCtx.Messages = append(deps.agentCtx.Messages, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: content, + }) + cfg := runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: deps.live.Model, + Provider: deps.live.ProviderName, + ThinkingLevel: deps.live.ThinkingLevel, + Stream: provider.StreamFnFromProvider(deps.live.Provider), + GetAPIKey: deps.creds.GetAPIKey, + ContextWindow: deps.live.ContextWindow, + Compaction: compaction.DefaultCompactionSettings, + }, + Batch: agenttool.BatchConfig{ + ToolExecutorConfig: agenttool.ToolExecutorConfig{ + Registry: deps.reg, + BeforeToolCall: beforeToolCall(deps, out), + }, + }, + Reminders: deps.reminders, + SessionID: deps.header.ID, + MemoryRoot: deps.memoryRoot, + } + // Per-turn wiring of the tool-execution + Stop seams (PreToolUse/PostToolUse/ + // Stop) onto this turn's freshly-built cfg; a nil dispatcher is a no-op so the + // hot path pays nothing when no hooks are configured (FR-18). + if deps.dispatcher != nil { + run.InstallSeams(&cfg, deps.dispatcher, deps.hookDeps) + } + stream := runtime.StartRun(ctx, deps.agentCtx, cfg) + + // The assistant reply is Markdown, which can only be laid out once the whole + // block is known, so streamed text is buffered here and rendered at turn end + // (ui.RenderMarkdown). On non-terminal output it returns the raw + // source, so pipes/tests are unchanged. flushReply guarantees the rendered + // block ends on a fresh line so tool activity below it starts cleanly. + var reply strings.Builder + flushReply := func() { + if reply.Len() == 0 { + return + } + rendered := ui.RenderMarkdown(reply.String()) + fmt.Fprint(out, rendered) + if !strings.HasSuffix(rendered, "\n") { + fmt.Fprintln(out) + } + reply.Reset() + } + // The SessionEnd/PreCompact observer chains onto the existing OnEvent closure + // (plugin notifier + telemetry). Nil when no hooks are configured. + var hookEvent func(agentcore.AgentEvent) + if deps.dispatcher != nil { + hookEvent = hooks.NewHookNotifier(deps.dispatcher, deps.hookDeps.SessionID, deps.hookDeps.ProjectDir).Handle + } + _, err = runtime.DrainStream(ctx, stream, runtime.StreamHandler{ + OnEvent: func(ev agentcore.AgentEvent) { + // First call the notifier if present. + if notifier := deps.notifierHandle(); notifier != nil { + notifier(ev) + } + // Capture TelemetryEvent and fold it into the holder. + if telemetry, ok := ev.(agentcore.TelemetryEvent); ok && deps.telemetry != nil { + deps.telemetry.Fold(telemetry) + } + // Deliver SessionEnd/PreCompact to any configured hook. + if hookEvent != nil { + hookEvent(ev) + } + // Surface auto-compaction progress: a "Compacting conversation…" line + // when it starts, then a token summary when it finishes. Manual + // /compact prints its own summary via runManualCompact (it bypasses + // the loop), so only threshold/overflow events reach here. + switch e := ev.(type) { + case agentcore.CompactionStartEvent: + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, "Compacting conversation…")) + case agentcore.CompactionEvent: + if e.ErrorMessage != "" { + fmt.Fprintf(out, "compaction failed: %s\n", e.ErrorMessage) + } else { + fmt.Fprintf(out, "compacted: %d → %d tokens, summarized %d messages, kept %d\n", + e.TokensBefore, e.TokensAfter, e.SummarizedCount, e.KeptCount) + } + } + }, + OnText: func(delta string) { + reply.WriteString(delta) + }, + OnTurnEnd: func(msg agentcore.AssistantMessage, results []agentcore.ToolResultMessage) { + flushReply() + for _, c := range msg.ToolCalls() { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(ui.Enabled(), ui.Green, "→ tool:"), ui.ToolCallLabel(c)) + } + for _, tr := range results { + ui.RenderToolResult(out, tr) + } + // Surface a failed turn so a provider/API error is never silent. The + // loop delivers request failures as a terminal assistant message with + // stopReason error/aborted (they do not ride the run's result error), + // so without this the REPL would print nothing at all for, e.g., a 4xx + // from the endpoint or an empty/unparseable response. + switch msg.StopReason { + case agentcore.StopReasonError: + reason := strings.TrimSpace(msg.ErrorMessage) + if reason == "" { + reason = "the provider returned an error with no message" + } + fmt.Fprintf(out, "%s %s\n", ui.Colorize(ui.Enabled(), ui.Red, "error:"), reason) + case agentcore.StopReasonAborted: + fmt.Fprintf(out, "%s aborted\n", ui.Colorize(ui.Enabled(), ui.Red, "error:")) + default: + // A turn that ends cleanly (end_turn) but produced no text, no + // thinking, and no tool calls means the endpoint accepted the + // request but sent back nothing usable (e.g. a 200 whose body was + // not in the wire format this protocol expects). Say so instead of + // returning to the prompt with no output. + if len(msg.Content) == 0 && len(results) == 0 { + fmt.Fprintf(out, "%s empty response from the model (no content). "+ + "Check that --model, --base-url and --protocol match the same provider.\n", + ui.Colorize(ui.Enabled(), ui.Yellow, "note:")) + } + } + }, + }) + // A run can end (error or interrupt) with buffered text from a final turn + // that never fired OnTurnEnd; flush it so no reply is silently dropped. + flushReply() + if err != nil { + if ctx.Err() != nil { + fmt.Fprintln(out, "^C interrupted") + } else { + fmt.Fprintf(out, "error: %v\n", err) + } + } +} + +// runForkClone handles the /fork and /clone commands (US-006, #122), which +// branch the conversation tree into a new, independent session. +// +// - "/clone" duplicates the entire current conversation at its current leaf +// into a fresh session and switches to it. Appending to the clone never +// affects the original (they are separate files). +// - "/fork" with no argument lists the historical user messages, numbered, so +// the user can pick one. +// - "/fork N" branches from BEFORE the N-th listed user message: the new +// session holds everything up to (but excluding) that message, so the user +// re-prompts from that point on an independent branch. +// +// Both first persist the current session (so the fork copies a saved tree), +// then call store.Fork to write the new branch, and finally swap deps.header and +// deps.agentCtx to the new session in place so subsequent prompts continue on +// the branch. resume of either session id later walks to its own leaf. +func runForkClone(out io.Writer, deps *replDeps, line string) { + // Persist the current turn as a branch so Fork copies an up-to-date tree + // without flattening any existing branches (a plain Save would rewrite the + // file linearly and drop siblings). + cli.PersistTurn(out, deps) + _, entries, err := deps.store.LoadEntries(deps.header.ID) + if err != nil && !errors.Is(err, os.ErrNotExist) { + fmt.Fprintf(out, "pigo: cannot read session tree: %v\n", err) + return + } + if len(entries) == 0 { + fmt.Fprintln(out, "nothing to fork yet — send a message first") + return + } + + fields := strings.Fields(line) + cmd := fields[0] + + var leafID string + if cmd == "/clone" { + // Clone the whole conversation at the current active leaf. + leafID = deps.curLeaf + if leafID == "" { + leafID = entries[len(entries)-1].ID + } + } else { // /fork + // Collect the historical user messages, in order, with their entry index. + type userMsg struct { + idx int // index into entries + text string + } + var users []userMsg + for i, e := range entries { + if u, ok := e.Message.(agentcore.UserMessage); ok { + users = append(users, userMsg{idx: i, text: agentcore.ContentToText(u.Content)}) + } + } + if len(users) == 0 { + fmt.Fprintln(out, "no user messages to fork from") + return + } + if len(fields) < 2 { + // List the user messages for selection. + fmt.Fprintln(out, "fork from which message? run /fork :") + for n, u := range users { + fmt.Fprintf(out, " %d. %s\n", n+1, ui.OneLine(u.text)) + } + return + } + n, convErr := strconv.Atoi(fields[1]) + if convErr != nil || n < 1 || n > len(users) { + fmt.Fprintf(out, "invalid selection %q — run /fork to list messages (1..%d)\n", fields[1], len(users)) + return + } + // Fork BEFORE the chosen user message: its parent becomes the new leaf, so + // the copied branch excludes that message and everything after it. + leafID = entries[users[n-1].idx].ParentID + } + + newHeader, path, err := deps.store.Fork(deps.header.ID, leafID, time.Now().UTC()) + if err != nil { + fmt.Fprintf(out, "pigo: fork failed: %v\n", err) + return + } + + // Swap the live session to the new branch: rebuild the flat message list from + // the copied path and point the REPL at the new header. Subsequent prompts + // grow the branch, never the original. + msgs := make(agentcore.MessageList, len(path)) + for i, e := range path { + msgs[i] = e.Message + } + deps.header = newHeader + deps.agentCtx.Messages = msgs + // The new file already holds the copied path verbatim, so mark it all + // persisted and set the active leaf to the copied tip. + deps.persisted = len(path) + deps.curLeaf = "" + if len(path) > 0 { + deps.curLeaf = path[len(path)-1].ID + } + // A side thread branched from the old conversation is meaningless against the + // new branch, so drop it (#281). + deps.lastBtw = nil + deps.lastBtwBase = 0 + // Reset telemetry for the new session, since cumulative stats should not bleed + // across conversations. + if deps.telemetry != nil { + deps.telemetry.Reset() + } + + action := "cloned" + if cmd == "/fork" { + action = "forked" + } + fmt.Fprintf(out, "%s session %s → %s (%d messages)\n", action, newHeader.ParentSession, newHeader.ID, len(msgs)) +} + +// runTree handles the /tree command (US-007, #123): the branch-navigation view +// for a pure line REPL. +// +// - "/tree" with no argument persists the current turn, then prints the +// session's entry tree with ├─/└─ connectors, tagging the active leaf with +// "← current". Each printed row is numbered so it can be selected. +// - "/tree N" switches the active leaf to the N-th printed entry: it rebuilds +// the shared context to the root→leaf path feeding that entry, so the next +// prompt continues from there. Because persistTurn appends new turns as a +// branch descending from the (moved) leaf, continuing after a switch grows a +// real sibling branch on disk rather than truncating history. +func runTree(out io.Writer, deps *replDeps, line string) { + // Persist any un-saved turn first so the tree reflects the live conversation. + cli.PersistTurn(out, deps) + _, entries, err := deps.store.LoadEntries(deps.header.ID) + if err != nil && !errors.Is(err, os.ErrNotExist) { + fmt.Fprintf(out, "pigo: cannot read session tree: %v\n", err) + return + } + if len(entries) == 0 { + fmt.Fprintln(out, "session tree is empty — send a message first") + return + } + + lines := session.RenderTreeLines(entries, deps.curLeaf) + fields := strings.Fields(line) + if len(fields) < 2 { + // Print the numbered tree for selection. + fmt.Fprintln(out, "session tree (run /tree to switch the active branch):") + for i, l := range lines { + fmt.Fprintf(out, " %d. %s\n", i+1, l.Text) + } + return + } + + n, convErr := strconv.Atoi(fields[1]) + if convErr != nil || n < 1 || n > len(lines) { + fmt.Fprintf(out, "invalid selection %q — run /tree to list nodes (1..%d)\n", fields[1], len(lines)) + return + } + + // Switch the active leaf to the chosen node and rebuild the context from its + // root→leaf path. The chosen entry is already persisted, so persisted stays at + // the rebuilt length; the next turn branches from curLeaf. + target := lines[n-1].Entry + path := session.PathToLeaf(entries, target.ID) + msgs := make(agentcore.MessageList, len(path)) + for i, e := range path { + msgs[i] = e.Message + } + deps.agentCtx.Messages = msgs + deps.curLeaf = target.ID + deps.persisted = len(msgs) + fmt.Fprintf(out, "switched to branch at node %d (%d messages) — next prompt continues from here\n", n, len(msgs)) +} + +// pathCommandArg extracts the path argument for a "/cmd " line, enforcing a +// command-token boundary (so "/exporter x" is NOT "/export" with arg "x") and +// stripping a single layer of surrounding double quotes (so a path with spaces +// can be given as /export "my session.html"). It returns "" when line is not the +// given command or carries no argument. +func pathCommandArg(line, cmd string) string { + if line != cmd && !strings.HasPrefix(line, cmd+" ") { + return "" + } + arg := strings.TrimSpace(strings.TrimPrefix(line, cmd)) + if len(arg) >= 2 && strings.HasPrefix(arg, `"`) && strings.HasSuffix(arg, `"`) { + arg = arg[1 : len(arg)-1] + } + return arg +} + +// runExport handles the /export command (US-008, #124): it persists the live +// turn, then writes the current session to a file. "/export" with no path +// defaults to ".jsonl" in the working directory; "/export path.html" +// (or .htm) writes a self-contained HTML transcript; any other extension writes +// JSONL. The JSONL form round-trips losslessly through /import. +func runExport(out io.Writer, deps *replDeps, line string) { + cli.PersistTurn(out, deps) + path := pathCommandArg(line, "/export") + if path == "" { + path = deps.header.ID + ".jsonl" + } + n, err := deps.store.Export(deps.header.ID, path) + if err != nil { + fmt.Fprintf(out, "pigo: export failed: %v\n", err) + return + } + fmt.Fprintf(out, "exported %d entries to %s\n", n, path) +} + +// runImport handles the /import command (US-008, #124): it reads a JSONL export +// and materializes it as a fresh, independent session, then switches the live +// REPL to it (swapping header + shared context in place) so the next prompt +// continues the imported conversation. The original file is untouched; the new +// session records the source id as its ParentSession. +func runImport(out io.Writer, deps *replDeps, line string) { + path := pathCommandArg(line, "/import") + if path == "" { + fmt.Fprintln(out, "usage: /import ") + return + } + newHeader, entries, err := deps.store.Import(path, time.Now().UTC()) + if err != nil { + fmt.Fprintf(out, "pigo: import failed: %v\n", err) + return + } + // Swap the live session to the imported one: rebuild the flat message list and + // point the REPL at the new header. The imported file already holds the entries + // verbatim, so mark them all persisted and set the active leaf to the tip. + msgs := make(agentcore.MessageList, len(entries)) + for i, e := range entries { + msgs[i] = e.Message + } + deps.header = newHeader + deps.agentCtx.Messages = msgs + deps.persisted = len(entries) + deps.curLeaf = "" + if len(entries) > 0 { + deps.curLeaf = entries[len(entries)-1].ID + } + // A side thread branched from the previous conversation no longer applies to + // the imported one, so drop it (#281). + deps.lastBtw = nil + deps.lastBtwBase = 0 + // Reset telemetry for the imported session, since cumulative stats should not + // bleed across conversations. + if deps.telemetry != nil { + deps.telemetry.Reset() + } + fmt.Fprintf(out, "imported %d entries from %s → session %s\n", len(entries), path, newHeader.ID) +} + +// runCopy handles the /copy command (US-009, #125): it copies the most recent +// assistant text message to the system clipboard. When no clipboard utility is +// available it degrades to printing the text with a notice, so the content is +// never lost. An empty conversation (no assistant reply yet) is reported. +func runCopy(out io.Writer, deps *replDeps) { + text := "" + for i := len(deps.agentCtx.Messages) - 1; i >= 0; i-- { + if a, ok := deps.agentCtx.Messages[i].(agentcore.AssistantMessage); ok { + if t := strings.TrimSpace(agentcore.ContentToText(a.Content)); t != "" { + text = t + break + } + } + } + if text == "" { + fmt.Fprintln(out, "nothing to copy — no assistant reply yet") + return + } + if err := clipboard.Copy(text); err != nil { + if errors.Is(err, clipboard.ErrUnavailable) { + // Degrade gracefully: print the text and tell the user why it was not + // copied, so the content is still recoverable. + fmt.Fprintln(out, "no clipboard utility found (install pbcopy/wl-copy/xclip/xsel); printing instead:") + fmt.Fprintln(out, text) + return + } + fmt.Fprintf(out, "pigo: copy failed: %v\n", err) + return + } + fmt.Fprintf(out, "copied last reply to clipboard (%d chars)\n", len(text)) +} + +// runSession handles the /session command (US-009, #125): it prints a summary of +// the live session — id, message count, estimated token usage, model/provider, +// creation time, and how many compaction checkpoints it contains. Counts are +// derived from the in-memory context (the source of truth for the live turn) so +// the numbers reflect unsaved messages too. +func runSession(out io.Writer, deps *replDeps) { + msgs := deps.agentCtx.Messages + tokens := compaction.EstimateContextTokens(msgs).Tokens + compactions := 0 + for _, m := range msgs { + if _, ok := m.(agentcore.CompactionMessage); ok { + compactions++ + } + } + fmt.Fprintf(out, "session: %s\n", deps.header.ID) + fmt.Fprintf(out, "messages: %d\n", len(msgs)) + fmt.Fprintf(out, "tokens (est): %d\n", tokens) + model := deps.live.Model + providerName := deps.live.ProviderName + if model == "" { + model = deps.header.Model + } + if providerName == "" { + providerName = deps.header.Provider + } + fmt.Fprintf(out, "model: %s (provider: %s)\n", model, providerName) + if !deps.header.CreatedAt.IsZero() { + fmt.Fprintf(out, "created: %s\n", deps.header.CreatedAt.Format(time.RFC3339)) + } + fmt.Fprintf(out, "compactions: %d\n", compactions) +} + +// runManualCompact compacts the shared context on an explicit /compact request: +// it runs the summarization stream, replaces the context with the checkpoint + +// retained tail, and prints the before/after token counts and retained message +// count. A failure is reported but non-fatal — the original context is kept +// unchanged (US-004). It uses the same provider/model as the live run. +func runManualCompact(out io.Writer, deps replDeps) { + msgs := deps.agentCtx.Messages + settings := compaction.DefaultCompactionSettings + before := compaction.EstimateContextTokens(msgs).Tokens + + stream := provider.StreamFnFromProvider(deps.live.Provider) + model := provider.Model{Provider: deps.live.ProviderName, ID: deps.live.Model, ContextWindow: deps.live.ContextWindow} + + // Resolve the API key like a normal turn so summarization authenticates + // against auth-requiring providers (otherwise Compact fails with + // "missing API key" and /compact would always report a non-fatal failure). + scfg := provider.StreamConfig{} + if deps.creds != nil { + scfg.APIKey = deps.creds.GetAPIKey(context.Background(), deps.live.ProviderName) + } + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, "Compacting conversation…")) + res, err := compaction.Compact(context.Background(), stream, model, msgs, settings, -1, nil, "", scfg) + if err != nil { + fmt.Fprintf(out, "compaction failed: %v (context left unchanged)\n", err) + return + } + if res == nil { + fmt.Fprintf(out, "nothing to compact (%d tokens, %d messages)\n", before, len(msgs)) + return + } + now := time.Now().UnixMilli() + rebuilt := res.RebuildContext(msgs, now) + deps.agentCtx.Messages = rebuilt + after := compaction.EstimateContextTokens(rebuilt).Tokens + summarized := len(msgs) - (len(rebuilt) - 1) + fmt.Fprintf(out, "compacted: %d → %d tokens, summarized %d messages, kept %d\n", + before, after, summarized, len(rebuilt)-1) +} + +// runManualRebuild reconstructs the shared context on an explicit /rebuild +// request. It reloads the session's persisted checkpoint (from #480) and inserts +// the compression boundary at its watermark — collapsing the pre-watermark +// prefix into the checkpoint summary and preserving the recent tail verbatim. If +// no checkpoint exists it falls back to the same lossy compaction /compact runs. +// It prints the before/after token counts; a failure is reported but non-fatal +// (the original context is kept unchanged). +func runManualRebuild(out io.Writer, deps replDeps) { + msgs := deps.agentCtx.Messages + before := compaction.EstimateContextTokens(msgs).Tokens + + // Build a RunConfig matching a normal turn so the no-checkpoint fallback can + // summarize against the live provider/model (RebuildFromCheckpoint reuses the + // loop's compaction path). The checkpoint path itself is pure/local. + cfg := runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: deps.live.Model, + Provider: deps.live.ProviderName, + ThinkingLevel: deps.live.ThinkingLevel, + Stream: provider.StreamFnFromProvider(deps.live.Provider), + ContextWindow: deps.live.ContextWindow, + Compaction: compaction.DefaultCompactionSettings, + }, + } + if deps.creds != nil { + cfg.GetAPIKey = deps.creds.GetAPIKey + } + + // Checkpoints live at /sessions//checkpoint.md; recover from + // the same root streamRun writes to. Empty when memory is disabled — then + // RebuildFromCheckpoint falls back to lossy compaction. + memoryRoot := deps.memoryRoot + + fmt.Fprintln(out, ui.Colorize(ui.Enabled(), ui.Dim, "Preparing conversation context…")) + res, err := runtime.RebuildFromCheckpoint(context.Background(), msgs, deps.header.ID, memoryRoot, &cfg, nil) + if err != nil { + fmt.Fprintf(out, "rebuild failed: %v (context left unchanged)\n", err) + return + } + if res.NoOp { + fmt.Fprintf(out, "nothing to rebuild (%d tokens, %d messages)\n", before, len(msgs)) + return + } + deps.agentCtx.Messages = res.Messages + source := "checkpoint" + if !res.FromCheckpoint { + source = "compaction (no checkpoint)" + } + fmt.Fprintf(out, "rebuilt from %s: %d → %d tokens, collapsed %d messages, kept %d\n", + source, res.TokensBefore, res.TokensAfter, res.SummarizedCount, res.KeptCount) +} + +// replayTranscript prints a resumed session's prior messages to out so the user +// sees the conversation so far before the first new prompt. +func replayTranscript(out io.Writer, messages []agentcore.AgentMessage) { + color := ui.Enabled() + for _, m := range messages { + switch msg := m.(type) { + case agentcore.UserMessage: + if t := agentcore.ContentToText(msg.Content); t != "" { + fmt.Fprintf(out, "> %s\n", t) + } + case agentcore.AssistantMessage: + if t := agentcore.ContentToText(msg.Content); t != "" { + fmt.Fprintln(out, t) + } + for _, c := range msg.ToolCalls() { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Green, "→ tool:"), ui.ToolCallLabel(c)) + } + case agentcore.ToolResultMessage: + if msg.IsError { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Red, "← error:"), ui.OneLine(agentcore.ContentToText(msg.Content))) + } else { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Green, "← result:"), ui.OneLine(agentcore.ContentToText(msg.Content))) + } + } + } +} + +// (Compact tool-activity renderers moved to internal/cli/ui: ui.RenderToolResult, +// ui.ToolCallLabel, ui.OneLine — shared by the REPL, /btw and /goal.) diff --git a/pigo/internal/cli/repl/repl_test.go b/pigo/internal/cli/repl/repl_test.go new file mode 100644 index 0000000..dd2b3b4 --- /dev/null +++ b/pigo/internal/cli/repl/repl_test.go @@ -0,0 +1,656 @@ +package repl + +// Tests for the line-based REPL (#106): slash-command dispatch (action prints +// message and does NOT run; prompt runs; unknown command errors and does NOT +// run; /exit and EOF exit cleanly), multi-turn history accumulation, and +// streaming assistant text. The REPL is driven with a fake provider so no +// network is involved — the whole read → run → stream-print loop runs over an +// in-memory input reader and output buffer. + +import ( + "bytes" + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" +) + +// replProvider is a minimal Provider that streams one scripted text turn per +// StreamCompletion call and records how many times it was called, so a test can +// assert whether a run was launched. +type replProvider struct { + reply string + calls int +} + +func (p *replProvider) Name() string { return "faux" } +func (p *replProvider) Models() []provider.Model { + return []provider.Model{{Provider: "faux", ID: "faux"}} +} + +func (p *replProvider) StreamCompletion(ctx context.Context, req provider.CompletionRequest) (*provider.AssistantMessageEventStream, error) { + p.calls++ + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + withText := partial + withText.Content = agentcore.ContentList{agentcore.NewTextContent(p.reply)} + final := withText + final.StopReason = agentcore.StopReasonEndTurn + s := provider.NewAssistantMessageEventStream(0) + go func() { + _ = s.Emit(ctx, provider.StreamStartEvent{Partial: partial}) + _ = s.Emit(ctx, provider.StreamTextEvent{Partial: withText}) + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: final}) + s.Close() + }() + return s, nil +} + +// newTestDeps builds replDeps wired to the fake provider and a temp session +// store, with a registry carrying one action command and one prompt command so +// slash dispatch can be exercised. actionRuns/promptResolved report whether each +// command fired. +func newTestDeps(t *testing.T, p provider.Provider) (replDeps, *session.Store) { + t.Helper() + store, err := session.NewStore(t.TempDir()) + if err != nil { + t.Fatalf("new store: %v", err) + } + live := &cli.LiveConfig{Model: "faux", ProviderName: "faux", Provider: p} + reg := runtime.NewSlashRegistry() + reg.AddBuiltin(runtime.SlashCommand{ + Name: "ping", + Action: func(string) string { return "pong" }, + }) + reg.AddUser(runtime.SlashCommand{ + Name: "echo", + Expand: func(args string) string { return "expanded: " + args }, + }) + deps := replDeps{ + store: store, + header: session.SessionHeader{ID: session.NewID(time.Now().UTC()), Model: "faux", Provider: "faux"}, + agentCtx: &agentcore.AgentContext{}, + live: live, + reg: agenttool.NewToolRegistry(), + slash: reg, + creds: provider.NewCredentialStore(nil), + } + return deps, store +} + +// TestREPLExitCommand verifies /exit ends the loop cleanly with no error and no +// agent run. +func TestREPLExitCommand(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL returned error: %v", err) + } + if p.calls != 0 { + t.Errorf("/exit must not launch a run, got %d calls", p.calls) + } +} + +// TestREPLQuitCommand verifies /quit is an alias for /exit. +func TestREPLQuitCommand(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/quit\n"), &out, deps); err != nil { + t.Fatalf("runREPL returned error: %v", err) + } + if p.calls != 0 { + t.Errorf("/quit must not launch a run, got %d calls", p.calls) + } +} + +// TestREPLEOFExits verifies EOF (no /exit) ends the loop cleanly. +func TestREPLEOFExits(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + // Empty input → immediate EOF. + if err := runREPL(strings.NewReader(""), &out, deps); err != nil { + t.Fatalf("EOF should exit cleanly, got: %v", err) + } + if p.calls != 0 { + t.Errorf("EOF with no input must not run, got %d calls", p.calls) + } +} + +// TestREPLFinalLineNoNewline verifies the ReadString-based loop handles a final +// input line without a trailing newline: the line is still run as a prompt and +// the loop then exits cleanly on EOF. (The previous bufio.Scanner had the same +// behavior; this pins it for the reader-based loop.) +func TestREPLFinalLineNoNewline(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("hello"), &out, deps); err != nil { + t.Fatalf("runREPL on no-trailing-newline input: %v", err) + } + if p.calls != 1 { + t.Errorf("runs fired = %d, want 1 (the final line should run once)", p.calls) + } +} + +// TestREPLEmptyLineIgnored verifies blank lines are skipped without running. +func TestREPLEmptyLineIgnored(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("\n \n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("blank lines must not run, got %d calls", p.calls) + } +} + +// TestREPLActionCommandNoRun verifies an action slash command prints its message +// and does NOT launch an agent run. +func TestREPLActionCommandNoRun(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/ping\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("action command must not run, got %d calls", p.calls) + } + if !strings.Contains(out.String(), "pong") { + t.Errorf("action command message not printed, out=%q", out.String()) + } +} + +// TestREPLPromptCommandRuns verifies a prompt slash command expands and launches +// a run. +func TestREPLPromptCommandRuns(t *testing.T) { + p := &replProvider{reply: "ack"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/echo hello\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 1 { + t.Fatalf("prompt command should launch exactly 1 run, got %d", p.calls) + } + // The expanded prompt must have been appended as the user message. + if len(deps.agentCtx.Messages) == 0 { + t.Fatal("expected messages in context after run") + } + first, ok := deps.agentCtx.Messages[0].(agentcore.UserMessage) + if !ok || agentcore.ContentToText(first.Content) != "expanded: hello" { + t.Errorf("first message = %T %q, want expanded prompt", deps.agentCtx.Messages[0], agentcore.ContentToText(first.Content)) + } +} + +// TestREPLUnknownCommandNoRun verifies an unknown slash command prints an error +// and does NOT run or crash. +func TestREPLUnknownCommandNoRun(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/nope\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("unknown command must not run, got %d calls", p.calls) + } + if !strings.Contains(strings.ToLower(out.String()), "unknown") { + t.Errorf("expected an unknown-command error line, out=%q", out.String()) + } +} + +// TestREPLModelSwitchTakesEffect verifies the /model action command switches the +// live model mid-session (via registerLiveCommands + resolveProvider) without +// launching a run, and that the switch is reflected in live for the next turn. +func TestREPLModelSwitchTakesEffect(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + // Register the real live action commands (/model, /models, /help) against the + // same live config the REPL runs on, so /model mutates it. + prompts.RegisterLiveCommands(deps.slash, deps.live) + + var out bytes.Buffer + // /model with no arg reports the current model; /model switches to an + // Ollama preset (no API key required); /exit ends the loop. + in := strings.NewReader("/model\n/model ollama/llama3.3\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("/model actions must not launch a run, got %d calls", p.calls) + } + if deps.live.Model != "ollama/llama3.3" || deps.live.ProviderName != "ollama" { + t.Errorf("live not switched: model=%q provider=%q", deps.live.Model, deps.live.ProviderName) + } + s := out.String() + if !strings.Contains(s, "faux") { + t.Errorf("/model (no arg) should report the current model, out=%q", s) + } + if !strings.Contains(s, "ollama/llama3.3") { + t.Errorf("/model switch should confirm the new model, out=%q", s) + } +} + +// TestREPLPersistsModelIntoHeader verifies that after a run the session header +// records the live model/provider (US-006: a /model switch is persisted with the +// session), by reloading the saved session and inspecting its header. +func TestREPLPersistsModelIntoHeader(t *testing.T) { + p := &replProvider{reply: "ok"} + deps, store := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("hello\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + headers, err := store.List() + if err != nil { + t.Fatalf("store.List: %v", err) + } + if len(headers) != 1 { + t.Fatalf("expected 1 saved session, got %d", len(headers)) + } + if headers[0].Model != "faux" || headers[0].Provider != "faux" { + t.Errorf("header model/provider = %q/%q, want live faux/faux", headers[0].Model, headers[0].Provider) + } +} + +// TestRenderToolResultTodoFull verifies the todo tool's multi-line progress +// block is rendered in full under a "← todo:" header (US-011: REPL shows +// progress on update), while a non-todo result stays a one-line summary. +func TestRenderToolResultTodoFull(t *testing.T) { + var out bytes.Buffer + ui.RenderToolResult(&out, agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, ToolName: "todo", + Content: agentcore.ContentList{agentcore.NewTextContent("Todos:\n [x] a\n [ ] b\n(1/2 completed)")}, + }) + s := out.String() + for _, want := range []string{"← todo:", "[x] a", "[ ] b", "(1/2 completed)"} { + if !strings.Contains(s, want) { + t.Errorf("todo render missing %q, out=%q", want, s) + } + } + + out.Reset() + ui.RenderToolResult(&out, agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, ToolName: "bash", + Content: agentcore.ContentList{agentcore.NewTextContent("line1\nline2")}, + }) + if got := out.String(); !strings.Contains(got, "← result: line1") { + t.Errorf("non-todo render = %q, want a one-line summary", got) + } +} + +// TestReplayTranscriptRendersRoles verifies a resumed session's prior messages +// are echoed by role (user / assistant / tool result) before the first new +// prompt (US-006 acceptance: resumed conversation is replayed). +func TestReplayTranscriptRendersRoles(t *testing.T) { + msgs := []agentcore.AgentMessage{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("what is 2+2")}}, + agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{agentcore.NewTextContent("Let me compute."), agentcore.NewToolCallContent("c1", "calc", []byte(`{"expr":"2+2"}`))}, + }, + agentcore.ToolResultMessage{RoleField: agentcore.RoleToolResult, ToolCallID: "c1", ToolName: "calc", Content: agentcore.ContentList{agentcore.NewTextContent("4")}}, + } + var out bytes.Buffer + replayTranscript(&out, msgs) + s := out.String() + for _, want := range []string{"> what is 2+2", "Let me compute.", `→ tool: calc {"expr":"2+2"}`, "← result: 4"} { + if !strings.Contains(s, want) { + t.Errorf("replay missing %q, out=%q", want, s) + } + } +} + +// TestREPLTreePrintsAndSwitchesBranch drives the /tree command end to end over +// the REPL: after two turns, "/tree" prints a numbered tree with the current-leaf +// marker; "/tree 1" switches the active leaf to the first node, so the next +// prompt branches from there — leaving the original branch intact on disk (US-007, +// #123). +func TestREPLTreePrintsAndSwitchesBranch(t *testing.T) { + p := &replProvider{reply: "reply"} + deps, store := newTestDeps(t, p) + var out bytes.Buffer + // Two turns build a linear history (4 messages), then /tree lists it, /tree 1 + // switches to the first node, a new prompt branches, then /exit. + in := strings.NewReader("first\nsecond\n/tree\n/tree 1\nbranched\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + s := out.String() + if !strings.Contains(s, "← current") { + t.Errorf("/tree should mark the current leaf, out=%q", s) + } + if !strings.Contains(s, "1. user:") { + t.Errorf("/tree should number entries starting at the root user message, out=%q", s) + } + if !strings.Contains(s, "switched to branch at node 1") { + t.Errorf("/tree 1 should confirm the switch, out=%q", s) + } + // The on-disk tree must retain both branches: the original 4-message line plus + // the new branch off node 1. Reload and confirm the root has 2 children. + _, entries, err := store.LoadEntries(deps.header.ID) + if err != nil { + t.Fatalf("LoadEntries: %v", err) + } + rootID := "" + for _, e := range entries { + if e.ParentID == "" { + rootID = e.ID + break + } + } + if rootID == "" { + t.Fatal("no root entry found") + } + kids := 0 + for _, e := range entries { + if e.ParentID == rootID { + kids++ + } + } + if kids != 2 { + t.Errorf("root should have 2 children after branching, got %d (entries=%d)", kids, len(entries)) + } +} + +// TestREPLTreeEmptyNoop verifies /tree on a fresh session (no messages) prints a +// friendly notice and does not crash or run. +func TestREPLTreeEmptyNoop(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/tree\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("/tree must not launch a run, got %d calls", p.calls) + } + if !strings.Contains(out.String(), "empty") { + t.Errorf("/tree on empty session should say so, out=%q", out.String()) + } +} + +// is printed, and history accumulates across two turns in the shared context. +func TestREPLStreamsAndAccumulatesHistory(t *testing.T) { + p := &replProvider{reply: "the answer"} + deps, store := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("first question\nsecond question\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 2 { + t.Fatalf("two prompts should launch 2 runs, got %d", p.calls) + } + // The streamed reply text must appear in the output. + if !strings.Contains(out.String(), "the answer") { + t.Errorf("assistant reply not streamed to output, out=%q", out.String()) + } + // History: user(first) + assistant + user(second) + assistant = 4 messages. + if len(deps.agentCtx.Messages) != 4 { + t.Fatalf("expected 4 accumulated messages, got %d", len(deps.agentCtx.Messages)) + } + u0, _ := deps.agentCtx.Messages[0].(agentcore.UserMessage) + u2, _ := deps.agentCtx.Messages[2].(agentcore.UserMessage) + if agentcore.ContentToText(u0.Content) != "first question" || agentcore.ContentToText(u2.Content) != "second question" { + t.Errorf("history not accumulated in order: %q, %q", agentcore.ContentToText(u0.Content), agentcore.ContentToText(u2.Content)) + } + // The session must have been persisted after the runs. + headers, err := store.List() + if err != nil { + t.Fatalf("store.List: %v", err) + } + if len(headers) != 1 { + t.Errorf("expected 1 saved session, got %d", len(headers)) + } +} + +// errProvider streams a single turn that ends with stopReason error carrying an +// ErrorMessage, mimicking how the loop surfaces a request failure (e.g. a 4xx +// from the endpoint) as a terminal assistant message rather than a Go error. +type errProvider struct { + reason string + calls int +} + +func (p *errProvider) Name() string { return "faux" } +func (p *errProvider) Models() []provider.Model { + return []provider.Model{{Provider: "faux", ID: "faux"}} +} + +func (p *errProvider) StreamCompletion(ctx context.Context, req provider.CompletionRequest) (*provider.AssistantMessageEventStream, error) { + p.calls++ + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + final := partial + final.StopReason = agentcore.StopReasonError + final.ErrorMessage = p.reason + s := provider.NewAssistantMessageEventStream(0) + go func() { + _ = s.Emit(ctx, provider.StreamStartEvent{Partial: partial}) + _ = s.Emit(ctx, provider.StreamErrorEvent{Message: final}) + s.Close() + }() + return s, nil +} + +// TestREPLSurfacesTurnError verifies a turn that ends with stopReason error is +// printed to the user instead of returning silently to the prompt. Without this +// an API failure (delivered as a terminal error message, not a run error) would +// produce no output at all. +func TestREPLSurfacesTurnError(t *testing.T) { + p := &errProvider{reason: "401 unauthorized: bad api key"} + deps, _ := newTestDeps(t, p) + deps.live.Provider = p + var out bytes.Buffer + if err := runREPL(strings.NewReader("hello\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 1 { + t.Fatalf("expected 1 run, got %d", p.calls) + } + got := out.String() + if !strings.Contains(got, "error:") || !strings.Contains(got, "401 unauthorized: bad api key") { + t.Errorf("turn error not surfaced to output, out=%q", got) + } +} + +// emptyProvider streams a clean end_turn with no content, thinking, or tool +// calls — the shape produced when an endpoint accepts the request with a 200 but +// returns nothing this protocol can decode. +type emptyProvider struct{ calls int } + +func (p *emptyProvider) Name() string { return "faux" } +func (p *emptyProvider) Models() []provider.Model { + return []provider.Model{{Provider: "faux", ID: "faux"}} +} + +func (p *emptyProvider) StreamCompletion(ctx context.Context, req provider.CompletionRequest) (*provider.AssistantMessageEventStream, error) { + p.calls++ + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + final := partial + final.StopReason = agentcore.StopReasonEndTurn + s := provider.NewAssistantMessageEventStream(0) + go func() { + _ = s.Emit(ctx, provider.StreamStartEvent{Partial: partial}) + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: final}) + s.Close() + }() + return s, nil +} + +// TestREPLNotesEmptyResponse verifies a clean turn that produced no output at +// all is flagged with a note rather than returning silently to the prompt. +func TestREPLNotesEmptyResponse(t *testing.T) { + p := &emptyProvider{} + deps, _ := newTestDeps(t, p) + deps.live.Provider = p + var out bytes.Buffer + if err := runREPL(strings.NewReader("hello\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 1 { + t.Fatalf("expected 1 run, got %d", p.calls) + } + if got := out.String(); !strings.Contains(got, "empty response from the model") { + t.Errorf("empty response not flagged, out=%q", got) + } +} + +// TestREPLExportImportRoundTrip drives /export and /import end to end over the +// REPL: after a turn, "/export " writes a JSONL file, then "/import +// " materializes it as a fresh session and switches to it (US-008, #124). +func TestREPLExportImportRoundTrip(t *testing.T) { + p := &replProvider{reply: "the answer"} + deps, store := newTestDeps(t, p) + origID := deps.header.ID + out := filepath.Join(t.TempDir(), "sess.jsonl") + + var buf bytes.Buffer + in := strings.NewReader("hello\n/export " + out + "\n/import " + out + "\n/exit\n") + if err := runREPL(in, &buf, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + s := buf.String() + if !strings.Contains(s, "exported") { + t.Errorf("/export should confirm, out=%q", s) + } + if !strings.Contains(s, "imported") { + t.Errorf("/import should confirm, out=%q", s) + } + // The export file must exist and be non-empty. + if info, err := os.Stat(out); err != nil || info.Size() == 0 { + t.Fatalf("export file missing or empty: err=%v", err) + } + // The import creates a new session distinct from the original, so the store + // should now hold at least 2 sessions. + headers, err := store.List() + if err != nil { + t.Fatalf("store.List: %v", err) + } + var foundNew bool + for _, h := range headers { + if h.ID != origID && h.ParentSession == origID { + foundNew = true + } + } + if !foundNew { + t.Errorf("expected an imported session with ParentSession=%q, headers=%+v", origID, headers) + } +} + +// TestREPLExportDefaultsToJSONL verifies "/export" with no path defaults to +// ".jsonl" and does not launch an agent run. +func TestREPLExportDefaultsToJSONL(t *testing.T) { + p := &replProvider{reply: "ok"} + deps, _ := newTestDeps(t, p) + dir := t.TempDir() + // Run inside a temp dir so the default relative filename lands there. + cwd, _ := os.Getwd() + if err := os.Chdir(dir); err != nil { + t.Fatalf("chdir: %v", err) + } + defer os.Chdir(cwd) + + var buf bytes.Buffer + if err := runREPL(strings.NewReader("hi\n/export\n/exit\n"), &buf, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + want := deps.header.ID + ".jsonl" + if _, err := os.Stat(filepath.Join(dir, want)); err != nil { + t.Errorf("default export file %q not created: %v", want, err) + } +} + +// TestREPLImportTokenBoundary verifies that a command sharing a prefix with +// /import (e.g. "/important") is NOT treated as /import — it falls through to +// slash resolution and reports an unknown command rather than importing. +func TestREPLImportTokenBoundary(t *testing.T) { + p := &replProvider{reply: "ok"} + deps, _ := newTestDeps(t, p) + var buf bytes.Buffer + if err := runREPL(strings.NewReader("/important\n/exit\n"), &buf, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if strings.Contains(buf.String(), "imported") { + t.Errorf("/important must not trigger /import, out=%q", buf.String()) + } + if p.calls != 0 { + t.Errorf("/important should not launch a run, got %d calls", p.calls) + } +} + +// TestREPLSessionStats drives the /session command: after a turn it prints the +// session id, message count, token estimate, model, and compaction count without +// launching another run (US-009, #125). +func TestREPLSessionStats(t *testing.T) { + p := &replProvider{reply: "the answer"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("hello\n/session\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 1 { + t.Errorf("/session must not launch a run (only the prompt), got %d calls", p.calls) + } + s := out.String() + for _, want := range []string{"session:", "messages:", "tokens (est):", "model:", "compactions:"} { + if !strings.Contains(s, want) { + t.Errorf("/session output missing %q, out=%q", want, s) + } + } + // After one turn the context holds user + assistant = 2 messages. + if !strings.Contains(s, "messages: 2") { + t.Errorf("/session should report 2 messages, out=%q", s) + } +} + +// TestREPLCopyEmpty verifies /copy on a session with no assistant reply prints a +// friendly notice rather than copying or crashing. +func TestREPLCopyEmpty(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("/copy\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("/copy must not launch a run, got %d calls", p.calls) + } + if !strings.Contains(out.String(), "nothing to copy") { + t.Errorf("/copy on empty session should say so, out=%q", out.String()) + } +} + +// TestREPLCopyDegradesToPrint verifies /copy degrades to printing the last reply +// when no clipboard utility is available (PATH pointed at an empty dir), so the +// content is never lost. +func TestREPLCopyDegradesToPrint(t *testing.T) { + t.Setenv("PATH", t.TempDir()) + p := &replProvider{reply: "the important answer"} + deps, _ := newTestDeps(t, p) + var out bytes.Buffer + if err := runREPL(strings.NewReader("ask\n/copy\n/exit\n"), &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + s := out.String() + if !strings.Contains(s, "no clipboard utility") { + t.Errorf("/copy should report missing clipboard utility, out=%q", s) + } + if !strings.Contains(s, "the important answer") { + t.Errorf("/copy should print the reply when degrading, out=%q", s) + } +} diff --git a/pigo/internal/cli/repl/rewind.go b/pigo/internal/cli/repl/rewind.go new file mode 100644 index 0000000..46299bd --- /dev/null +++ b/pigo/internal/cli/repl/rewind.go @@ -0,0 +1,160 @@ +// This file implements the /rewind command (edit checkpoint / rewind): pigo's +// analogue of Claude Code's Esc-Esc rewind. Where /tree only moves the +// conversation leaf, /rewind also restores the working tree — it replays the +// file-snapshot journal (see agenttool.FileSnapshotRecorder) so a turn's write +// and edit mutations are rolled back, then switches the active conversation leaf +// to the point before that turn. The two together return the session to an +// earlier state in code and dialogue at once. +// +// Scope (v1): only pigo's own write/edit tools are journaled. Files changed by +// bash commands are not captured and are left untouched by a rewind. +package repl + +import ( + "fmt" + "io" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/session" +) + +// rewindLabel derives a short one-line description of a turn from its prompt, for +// display in the /rewind list. It collapses whitespace and truncates so the list +// stays scannable. +func rewindLabel(prompt string) string { + label := strings.Join(strings.Fields(prompt), " ") + const max = 60 + if len(label) > max { + label = label[:max-1] + "…" + } + return label +} + +// runRewind handles the /rewind command. With no argument it persists the live +// turn and prints the numbered restore points (most useful last). With "/rewind +// N" it restores files to their state before the N-th listed point and switches +// the conversation to the leaf that preceded that turn. +func runRewind(out io.Writer, deps *replDeps, line string) { + if deps.snap == nil { + fmt.Fprintln(out, "rewind is unavailable (file tools are disabled)") + return + } + // Persist any un-saved turn first so the just-run turn's restore point exists + // and the leaf ids we switch to are on disk. + cli.PersistTurn(out, deps) + + points := deps.snap.Points() + fields := strings.Fields(line) + if len(fields) < 2 { + printRewindPoints(out, points) + return + } + if len(points) == 0 { + fmt.Fprintln(out, "no restore points yet — file edits create them") + return + } + + n, err := strconv.Atoi(fields[1]) + if err != nil || n < 1 || n > len(points) { + fmt.Fprintf(out, "invalid selection %q — run /rewind to list points (1..%d)\n", fields[1], len(points)) + return + } + + leafID, restored, warnings, rErr := deps.snap.Restore(n - 1) + if rErr != nil { + fmt.Fprintf(out, "pigo: rewind failed: %v\n", rErr) + return + } + + if len(restored) > 0 { + fmt.Fprintf(out, "restored %d file(s):\n", len(restored)) + for _, p := range restored { + fmt.Fprintf(out, " %s\n", displayPath(deps.cwd, p)) + } + } else { + fmt.Fprintln(out, "no files to restore for this point") + } + for _, w := range warnings { + fmt.Fprintf(out, " warning: %s\n", w) + } + + // Move the conversation back to the leaf that preceded the turn, rebuilding the + // shared context from that leaf's root→leaf path (same mechanism as /tree). An + // empty leaf id means the turn was the first in the session: reset to an empty + // conversation. + if !rewindConversation(out, deps, leafID) { + return + } + fmt.Fprintf(out, "rewound to before point %d — next prompt continues from here\n", n) +} + +// rewindConversation switches the active leaf to leafID and rebuilds the shared +// context from its path. A "" leafID resets to an empty conversation (the turn +// was the session's first). It reports whether the switch succeeded. +func rewindConversation(out io.Writer, deps *replDeps, leafID string) bool { + if leafID == "" { + deps.agentCtx.Messages = nil + deps.curLeaf = "" + deps.persisted = 0 + return true + } + _, entries, err := deps.store.LoadEntries(deps.header.ID) + if err != nil { + fmt.Fprintf(out, "pigo: cannot read session tree: %v\n", err) + return false + } + path := session.PathToLeaf(entries, leafID) + if len(path) == 0 { + fmt.Fprintf(out, "pigo: restore point's conversation node is no longer in the tree; files were restored but the conversation was left unchanged\n") + return false + } + msgs := make(agentcore.MessageList, len(path)) + for i, e := range path { + msgs[i] = e.Message + } + deps.agentCtx.Messages = msgs + deps.curLeaf = leafID + deps.persisted = len(msgs) + return true +} + +// printRewindPoints renders the numbered restore points, oldest first, showing +// when each was made, how many files it touched, and the turn's label. +func printRewindPoints(out io.Writer, points []agenttool.RestorePoint) { + if len(points) == 0 { + fmt.Fprintln(out, "no restore points yet — file edits create them") + return + } + fmt.Fprintln(out, "restore points (run /rewind to roll files + conversation back to before that point):") + for i, p := range points { + files := len(p.Snapshots) + unit := "files" + if files == 1 { + unit = "file" + } + when := p.Time.Local().Format(time.Kitchen) + label := p.Label + if label == "" { + label = "(no prompt)" + } + fmt.Fprintf(out, " %d. %s %d %s %s\n", i+1, when, files, unit, label) + } +} + +// displayPath shortens an absolute snapshot path to a workspace-relative form for +// display when it lives under cwd; otherwise it returns the absolute path. +func displayPath(cwd, abs string) string { + if cwd == "" { + return abs + } + if rel, err := filepath.Rel(cwd, abs); err == nil && !strings.HasPrefix(rel, "..") { + return rel + } + return abs +} diff --git a/pigo/internal/cli/repl/rewind_test.go b/pigo/internal/cli/repl/rewind_test.go new file mode 100644 index 0000000..cd3d889 --- /dev/null +++ b/pigo/internal/cli/repl/rewind_test.go @@ -0,0 +1,89 @@ +// Tests for the /rewind command wiring: listing restore points and restoring +// files + conversation. The file-snapshot journal itself is tested in +// agenttool; here we exercise runRewind's REPL-level behavior over a replDeps. +package repl + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" +) + +// /rewind with no argument lists the committed restore points. +func TestREPLRewindListsPoints(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + deps.snap = agenttool.NewFileSnapshotRecorder() + + dir := t.TempDir() + f := filepath.Join(dir, "a.txt") + if err := os.WriteFile(f, []byte("v0"), 0o644); err != nil { + t.Fatal(err) + } + deps.snap.Record(f) + deps.snap.Commit("", "add feature X") + + var out bytes.Buffer + runRewind(&out, &deps, "/rewind") + got := out.String() + if !strings.Contains(got, "restore points") || !strings.Contains(got, "add feature X") { + t.Errorf("listing missing expected content:\n%s", got) + } +} + +// /rewind N restores the file to its baseline and resets the conversation when +// the point's leaf is empty (it was the session's first turn). +func TestREPLRewindRestoresFiles(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + deps.snap = agenttool.NewFileSnapshotRecorder() + deps.agentCtx.Messages = agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser}, + } + deps.persisted = 1 + + dir := t.TempDir() + f := filepath.Join(dir, "a.txt") + if err := os.WriteFile(f, []byte("original"), 0o644); err != nil { + t.Fatal(err) + } + deps.snap.Record(f) // baseline "original" + if err := os.WriteFile(f, []byte("changed"), 0o644); err != nil { + t.Fatal(err) + } + deps.snap.Commit("", "edit a.txt") + + var out bytes.Buffer + runRewind(&out, &deps, "/rewind 1") + + if data, _ := os.ReadFile(f); string(data) != "original" { + t.Errorf("file not restored: got %q, want original", string(data)) + } + if len(deps.agentCtx.Messages) != 0 { + t.Errorf("conversation not reset: %d messages remain", len(deps.agentCtx.Messages)) + } + if len(deps.snap.Points()) != 0 { + t.Errorf("journal not truncated after rewind") + } + if !strings.Contains(out.String(), "rewound to before point 1") { + t.Errorf("missing confirmation:\n%s", out.String()) + } +} + +// /rewind is unavailable when file tools are disabled (nil recorder). +func TestREPLRewindDisabled(t *testing.T) { + p := &replProvider{reply: "hi"} + deps, _ := newTestDeps(t, p) + deps.snap = nil + + var out bytes.Buffer + runRewind(&out, &deps, "/rewind") + if !strings.Contains(out.String(), "unavailable") { + t.Errorf("want unavailable message, got:\n%s", out.String()) + } +} diff --git a/pigo/internal/cli/repl/skills_test.go b/pigo/internal/cli/repl/skills_test.go new file mode 100644 index 0000000..26ae41f --- /dev/null +++ b/pigo/internal/cli/repl/skills_test.go @@ -0,0 +1,165 @@ +package repl + +// Tests for default skill loading from ~/.agents/skills and its exposure as +// /skill-name slash commands. skillsDir honors the PIGO_SKILLS_DIR override so +// the loader can be pointed at a temp dir without touching the real home dir. + +import ( + "os" + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/runtime" +) + +// writeSkill creates a skill markdown file with the given frontmatter body. +func writeSkill(t *testing.T, dir, name, content string) { + t.Helper() + if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644); err != nil { + t.Fatalf("write skill %s: %v", name, err) + } +} + +// findSkill returns the skill with the given name from the set, or nil. +func findSkill(skills []*runtime.Skill, name string) *runtime.Skill { + for _, s := range skills { + if s.Frontmatter.Name == name { + return s + } + } + return nil +} + +// TestLoadSkillsFromDir verifies skills in PIGO_SKILLS_DIR are loaded and expose +// a /skill-name slash command whose expansion is the skill body. +func TestLoadSkillsFromDir(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", dir) + t.Setenv("PIGO_HOME", t.TempDir()) + writeSkill(t, dir, "greet.md", "---\nname: greet\ndescription: say hello\n---\nYou are a friendly greeter.") + + skills, err := run.LoadSkills(false) + if err != nil { + t.Fatalf("run.LoadSkills: %v", err) + } + s := findSkill(skills, "greet") + if s == nil { + t.Fatal("greet skill not loaded") + } + c := s.SlashCommand() + if c.Name != "greet" { + t.Errorf("Name = %q, want greet", c.Name) + } + if c.Description != "say hello" { + t.Errorf("Description = %q, want 'say hello'", c.Description) + } + if c.Expand == nil { + t.Fatal("skill command must be a prompt command (Expand != nil)") + } + if got := c.Expand(""); got != "You are a friendly greeter." { + t.Errorf("Expand(\"\") = %q, want the skill body", got) + } +} + +// TestLoadSkillsNoSkills verifies --no-skills skips discovery entirely: no +// skills are loaded and the skills dir is left untouched (no bootstrap). +func TestLoadSkillsNoSkills(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", dir) + t.Setenv("PIGO_HOME", t.TempDir()) + + skills, err := run.LoadSkills(true) + if err != nil { + t.Fatalf("run.LoadSkills(true): %v", err) + } + if len(skills) != 0 { + t.Errorf("got %d skills, want 0 under --no-skills", len(skills)) + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("read skills dir: %v", err) + } + if len(entries) != 0 { + t.Errorf("--no-skills must not bootstrap; skills dir has %d entries", len(entries)) + } +} + +// TestBuildSlashRegistryIncludesSkills verifies buildSlashRegistry wires the +// pre-loaded skills into the registry so /skill-name resolves. +func TestBuildSlashRegistryIncludesSkills(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", dir) + // Keep the user-commands path from touching a real home dir. + t.Setenv("PIGO_HOME", t.TempDir()) + writeSkill(t, dir, "summarize.md", "---\nname: summarize\ndescription: summarize input\n---\nSummarize the following: $ARGUMENTS") + + skills, err := run.LoadSkills(false) + if err != nil { + t.Fatalf("run.LoadSkills: %v", err) + } + reg, err := prompts.BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, skills, nil, prompts.PromptTemplateSources{}) + if err != nil { + t.Fatalf("buildSlashRegistry: %v", err) + } + out, err := reg.ResolveOutcome("/summarize hello world") + if err != nil { + t.Fatalf("ResolveOutcome: %v", err) + } + if !out.Handled { + t.Fatal("/summarize should be handled by the registry") + } + if out.Kind != runtime.SlashPrompt { + t.Errorf("Kind = %v, want SlashPrompt", out.Kind) + } + if out.Prompt != "Summarize the following: hello world" { + t.Errorf("Prompt = %q, want $ARGUMENTS substituted", out.Prompt) + } +} + +// TestBuildSlashRegistryNoSkills verifies that when no skills are passed (as +// under --no-skills, where loadSkills returns nil), a /skill-name command is not +// registered even though a skill file exists on disk (mirrors pi's --no-skills). +func TestBuildSlashRegistryNoSkills(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + + reg, err := prompts.BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, nil, nil, prompts.PromptTemplateSources{}) + if err != nil { + t.Fatalf("buildSlashRegistry: %v", err) + } + // With no skills registered, /summarize is an unknown command (an error) + // rather than a handled one. + if _, err := reg.ResolveOutcome("/summarize hello world"); err == nil { + t.Error("/summarize must be unknown when no skills are registered") + } +} + +// TestLoadSkillsBootstrapsBuiltinSkills verifies that on a fresh +// PIGO_SKILLS_DIR the built-in skills are installed and loaded (first-run +// bootstrap), so e.g. /prd resolves without any manual install. +func TestLoadSkillsBootstrapsBuiltinSkills(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", dir) + t.Setenv("PIGO_HOME", t.TempDir()) + + skills, err := run.LoadSkills(false) + if err != nil { + t.Fatalf("run.LoadSkills: %v", err) + } + reg, err := prompts.BuildSlashRegistry(&cli.LiveConfig{Model: "test", ProviderName: "test"}, skills, nil, prompts.PromptTemplateSources{}) + if err != nil { + t.Fatalf("buildSlashRegistry: %v", err) + } + for _, name := range []string{"/prd", "/refactor", "/architecture-diagram", "/weather"} { + out, err := reg.ResolveOutcome(name) + if err != nil { + t.Errorf("%s should be registered after bootstrap: %v", name, err) + continue + } + if !out.Handled { + t.Errorf("%s should be handled by the registry after bootstrap", name) + } + } +} diff --git a/pigo/internal/cli/repl/status_repl_test.go b/pigo/internal/cli/repl/status_repl_test.go new file mode 100644 index 0000000..77d4544 --- /dev/null +++ b/pigo/internal/cli/repl/status_repl_test.go @@ -0,0 +1,127 @@ +// This file holds the REPL-integration and headless-flag tests for /status +// (US-005, #295) that drive the package-main REPL harness (runREPL/newTestDeps) +// or inspect CLI flags. The direct-call rendering tests live in +// internal/cli/status alongside RunStatus. +package repl + +import ( + "bytes" + "strings" + "testing" + + flag "github.com/spf13/pflag" +) + +func TestStatusGuard(t *testing.T) { + // The guard logic is in the REPL loop: matches "/status" or "/status " + // but not "/statusfoo" or "/statusbar" + testCases := []struct { + line string + want bool + }{ + {"/status", true}, + {"/status ", true}, + {"/status foo", true}, + {"/statusbar", false}, + {"/statusfoo", false}, + {"/status123", false}, + {"/stat", false}, + {"/session", false}, + } + + for _, tc := range testCases { + got := (tc.line == "/status" || strings.HasPrefix(tc.line, "/status ")) + if got != tc.want { + t.Errorf("line %q: got %v, want %v", tc.line, got, tc.want) + } + } +} + +func TestRunStatusViaREPL(t *testing.T) { + // Verify that /status is intercepted in the REPL loop and doesn't invoke the model + p := &replProvider{reply: "should not be called"} + deps, _ := newTestDeps(t, p) + + var out bytes.Buffer + in := strings.NewReader("/status\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + + if p.calls != 0 { + t.Errorf("expected 0 model calls for /status, got %d", p.calls) + } + + output := out.String() + if !strings.Contains(output, "runtime config:") { + t.Error("expected REPL /status output to contain 'runtime config:'") + } + if !strings.Contains(output, "context:") { + t.Error("expected REPL /status output to contain 'context:'") + } +} + +func TestStatusFooNotIntercepted(t *testing.T) { + // Verify that "/statusfoo" is NOT intercepted as "/status" + p := &replProvider{reply: "model called"} + deps, _ := newTestDeps(t, p) + + var out bytes.Buffer + in := strings.NewReader("/statusfoo\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + + if p.calls != 0 { + t.Errorf("expected 0 model calls for /statusfoo, got %d", p.calls) + } + + output := out.String() + if strings.Contains(output, "runtime config:") { + t.Error("expected /statusfoo to NOT run the status command") + } +} + +// TestStatusNotInHeadless verifies /status is REPL-only: there is no --status +// CLI flag. (The /status intercept lives in runREPL only; headless print mode +// never runs the REPL loop, so "/status" there is treated as an ordinary +// prompt, not a command.) +func TestStatusNotInHeadless(t *testing.T) { + if f := flag.Lookup("status"); f != nil { + t.Errorf("--status flag should not exist (headless must not expose /status), got %v", f) + } +} + +// TestStatusE2EViaREPL drives /status through the REPL loop intercept and +// asserts every section is present and the model is never invoked. +func TestStatusE2EViaREPL(t *testing.T) { + p := &replProvider{reply: "should not be called"} + deps, _ := newTestDeps(t, p) + deps.cwd = "/tmp/e2e-repl" + deps.live.Model = "e2e-model" + deps.live.ProviderName = "e2e-prov" + deps.live.ContextWindow = 128000 + + var out bytes.Buffer + in := strings.NewReader("/status\n/exit\n") + if err := runREPL(in, &out, deps); err != nil { + t.Fatalf("runREPL: %v", err) + } + if p.calls != 0 { + t.Errorf("expected 0 model calls for /status, got %d", p.calls) + } + got := out.String() + for _, want := range []string{ + "runtime config:", + "model: e2e-model", + "context:", + "project & environment:", + "credentials & connectivity:", + "telemetry:", + "no telemetry yet", + } { + if !strings.Contains(got, want) { + t.Errorf("REPL /status: expected output to contain %q", want) + } + } +} diff --git a/pigo/internal/cli/run/hooks_converge_test.go b/pigo/internal/cli/run/hooks_converge_test.go new file mode 100644 index 0000000..bc37f0c --- /dev/null +++ b/pigo/internal/cli/run/hooks_converge_test.go @@ -0,0 +1,106 @@ +package run + +// Regression coverage for the #425 driver convergence (FR-16): the SAME resolved +// hook set must fire in every driver mode. Rather than stand up a provider-backed +// run for each of the six drivers, this pins the two DISTINCT wiring paths they +// route through — the one-shot path (headless / subagent_rpc via InstallDriverHooks +// / InstallHooks) and the multi-turn path (repl / tui / goal / btw via +// BuildDispatcher once + InstallSeams per turn) — and asserts a PreToolUse hook +// installed through either path reaches the BeforeToolCall seam and blocks. It +// also pins FR-18: an empty hook set wires no seam in either path, so a run with +// no hooks configured behaves exactly as before. + +import ( + "context" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// blockingPreToolUse is a hook set whose PreToolUse hook exits 2 (Claude Code +// block semantics), so any wired BeforeToolCall seam must return a blocking +// decision when it fires. +func blockingPreToolUse() hooks.HookSet { + return hooks.HookSet{ + "PreToolUse": {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: "exit 2"}}}}, + } +} + +// fireBeforeToolCall drives the wired BeforeToolCall seam once and reports +// whether it produced a blocking decision. A nil seam (no hook wired) reports +// false. +func fireBeforeToolCall(cfg *runtime.RunConfig) bool { + seam := cfg.Batch.ToolExecutorConfig.BeforeToolCall + if seam == nil { + return false + } + dec := seam(context.Background(), agentcore.AgentToolCall{Name: "Bash"}) + return dec != nil && dec.Block +} + +// TestHookConvergenceBothPaths asserts the same PreToolUse hook set fires in both +// driver wiring paths: the one-shot headless path and the multi-turn REPL path. +func TestHookConvergenceBothPaths(t *testing.T) { + deps := HookDeps{SessionID: "s1", ProjectDir: t.TempDir()} + set := blockingPreToolUse() + + // Headless / subagent_rpc path: InstallDriverHooks wires the seams all-in-one. + t.Run("headless", func(t *testing.T) { + var cfg runtime.RunConfig + d, _ := InstallDriverHooks(context.Background(), &cfg, set, deps, "startup", nil) + if d == nil { + t.Fatal("expected dispatcher for non-empty hook set") + } + if !fireBeforeToolCall(&cfg) { + t.Fatal("headless path: PreToolUse hook did not block") + } + }) + + // REPL / TUI / goal / btw path: BuildDispatcher once, then InstallSeams per turn. + t.Run("repl", func(t *testing.T) { + d := BuildDispatcher(set, deps) + if d == nil { + t.Fatal("expected dispatcher for non-empty hook set") + } + var cfg runtime.RunConfig + InstallSeams(&cfg, d, deps) + if !fireBeforeToolCall(&cfg) { + t.Fatal("repl path: PreToolUse hook did not block") + } + }) +} + +// TestHookConvergenceNoHooksUnchanged pins FR-18: with no hooks configured neither +// wiring path installs a BeforeToolCall seam, so both drivers behave exactly as +// they did before hooks existed. +func TestHookConvergenceNoHooksUnchanged(t *testing.T) { + deps := HookDeps{ProjectDir: t.TempDir()} + + t.Run("headless", func(t *testing.T) { + var cfg runtime.RunConfig + d, ev := InstallDriverHooks(context.Background(), &cfg, nil, deps, "startup", nil) + if d != nil { + t.Fatalf("expected nil dispatcher for empty hook set, got %v", d) + } + if ev != nil { + t.Fatal("expected event handler unchanged (nil) for empty hook set") + } + if cfg.Batch.ToolExecutorConfig.BeforeToolCall != nil { + t.Fatal("headless path: seam wired despite no hooks") + } + }) + + t.Run("repl", func(t *testing.T) { + d := BuildDispatcher(nil, deps) + if d != nil { + t.Fatalf("expected nil dispatcher for empty hook set, got %v", d) + } + var cfg runtime.RunConfig + InstallSeams(&cfg, d, deps) // nil dispatcher must be a no-op + if cfg.Batch.ToolExecutorConfig.BeforeToolCall != nil { + t.Fatal("repl path: seam wired despite no hooks") + } + }) +} diff --git a/pigo/internal/cli/run/hooks_driver.go b/pigo/internal/cli/run/hooks_driver.go new file mode 100644 index 0000000..975a0a0 --- /dev/null +++ b/pigo/internal/cli/run/hooks_driver.go @@ -0,0 +1,64 @@ +// This file provides the single convergence entry point every driver calls to +// wire hooks into its run (#425, FR-16). Before this, each of the six RunConfig +// assembly sites (repl/goal/btw/tui/headless/subagent_rpc) built the loop config +// independently, which was the main risk of a hook point being silently dropped +// in one mode but not another. Routing them all through InstallDriverHooks makes +// "which hook points a run has" a single decision rather than six. +// +// InstallDriverHooks resolves the trust-gated hook set (FR-14), installs the +// tool-execution and Stop seams via InstallHooks, dispatches SessionStart inline +// so injected context reaches turn one (#423), and chains the observer notifier +// (SessionEnd/PreCompact, #424) onto the driver's event seam. It returns the +// Dispatcher so a caller can additionally run UserPromptSubmit (prompt entry) or +// InstallSubagentStop (sub-agent), and the possibly-wrapped event handler to +// install on whatever OnEvent seam the driver owns (HeadlessConfig.OnEvent for +// headless, the DrainStream handler for the REPL/TUI). +package run + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// InstallDriverHooks is the uniform hook-wiring seam for every driver. It +// installs the tool-execution + Stop hook points onto cfg, dispatches +// SessionStart (registering any additionalContext as a one-shot reminder so it +// lands in turn one), and chains the SessionEnd/PreCompact observer onto onEvent. +// +// set is the already-resolved, trust-gated hook set (see ResolveHookSet). When +// it is empty NewDispatcher returns nil, so this is a no-op that returns +// (nil, onEvent) and the hot path pays nothing (FR-18): the run behaves exactly +// as it did before hooks existed. +// +// The returned dispatcher (nil when no hooks) lets the caller wire the remaining +// prompt-scoped / sub-agent hooks. The returned handler is onEvent unchanged when +// there are no hooks, or onEvent with the notifier chained after it otherwise — +// so the driver installs one handler regardless. +func InstallDriverHooks(ctx context.Context, cfg *runtime.RunConfig, set hooks.HookSet, deps HookDeps, source string, onEvent func(agentcore.AgentEvent)) (*hooks.Dispatcher, func(agentcore.AgentEvent)) { + d := InstallHooks(cfg, set, deps) + if d == nil { + return nil, onEvent + } + DispatchSessionStart(ctx, d, cfg, deps, source) + n := hooks.NewHookNotifier(d, deps.SessionID, deps.ProjectDir) + return d, chainEvent(onEvent, n.Handle) +} + +// chainEvent composes two AgentEvent observers into one that calls prev then +// next. A nil operand is identity, so chaining onto an unset seam returns the +// other unchanged (and returns nil when both are nil, keeping the seam unset). +func chainEvent(prev, next func(agentcore.AgentEvent)) func(agentcore.AgentEvent) { + if prev == nil { + return next + } + if next == nil { + return prev + } + return func(ev agentcore.AgentEvent) { + prev(ev) + next(ev) + } +} diff --git a/pigo/internal/cli/run/hooks_install.go b/pigo/internal/cli/run/hooks_install.go new file mode 100644 index 0000000..a56f705 --- /dev/null +++ b/pigo/internal/cli/run/hooks_install.go @@ -0,0 +1,223 @@ +// This file provides the single cli-layer assembly helper that composes hook +// dispatch into a runtime.RunConfig. It lives in the cli layer (not runtime) to +// avoid a runtime→hooks→runtime import cycle: runtime stays hook-agnostic and +// only exposes the generic seams (BeforeToolCall/AfterToolCall/ShouldStopAfterTurn), +// while this helper knows about both the resolved Config.Hooks and the seams. +// +// #419 establishes the skeleton: InstallHooks builds a Dispatcher (or short- +// circuits to nil when no hooks are configured, FR-18) and offers generic +// decorator combinators for the seams. The concrete per-event wiring is filled +// in by later issues (#420–#424); this file deliberately wires no event yet. +package run + +import ( + "context" + "encoding/json" + "io" + "os" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// HookDeps carries the run-scoped context a Dispatcher needs: the session id and +// project directory that populate each HookInput / the hook process environment, +// and the writer that receives isolation warnings. WarnLog may be nil (defaults +// to os.Stderr), matching the plugin.EventNotifier convention. +type HookDeps struct { + SessionID string + ProjectDir string + WarnLog io.Writer +} + +// InstallHooks builds the Dispatcher for a run from the resolved hook set and +// wires the tool-execution hook points into cfg. It short-circuits when no hooks +// are configured: NewDispatcher returns nil for an empty set, so the hot path +// pays nothing (FR-18) and nothing is wrapped. The returned dispatcher is used +// by later per-event wiring (#421–#424). +// +// PreToolUse is CHAINED onto the existing BeforeToolCall seam (occupied by the +// trust gate) rather than replacing it: trust runs first and stays authoritative +// (a trust block short-circuits before the user hook runs). PostToolUse is +// chained onto AfterToolCall as a last-writer so it can append feedback to an +// already-executed tool's result without undoing it. +func InstallHooks(cfg *runtime.RunConfig, set hooks.HookSet, deps HookDeps) *hooks.Dispatcher { + warn := deps.WarnLog + if warn == nil { + warn = os.Stderr + } + d := hooks.NewDispatcher(set, deps.ProjectDir, warn) + if d == nil { + return nil + } + InstallSeams(cfg, d, deps) + return d +} + +// BuildDispatcher builds the run's Dispatcher from the resolved hook set without +// touching a RunConfig. It is the entry point for the multi-turn drivers (REPL / +// TUI) that resolve hooks and fire SessionStart once per session, then install +// the per-turn seams (InstallSeams) on each turn's freshly-built RunConfig. It +// returns nil for an empty set (FR-18), so callers gate all hook work on non-nil. +func BuildDispatcher(set hooks.HookSet, deps HookDeps) *hooks.Dispatcher { + warn := deps.WarnLog + if warn == nil { + warn = os.Stderr + } + return hooks.NewDispatcher(set, deps.ProjectDir, warn) +} + +// InstallSeams wires the tool-execution and Stop hook points onto cfg from an +// already-built dispatcher. It is the shared body of InstallHooks and the +// per-turn install path for multi-turn drivers. A nil dispatcher is a no-op, so +// callers can invoke it unconditionally. +// +// PreToolUse is CHAINED onto the existing BeforeToolCall seam (occupied by the +// trust gate) rather than replacing it: trust runs first and stays authoritative +// (a trust block short-circuits before the user hook runs). PostToolUse is +// chained onto AfterToolCall as a last-writer so it can append feedback to an +// already-executed tool's result without undoing it. The Stop hook is chained +// onto the loop's natural-end seam (FR-10), bounded by the decorator's +// consecutive-block counter (FR-12). +func InstallSeams(cfg *runtime.RunConfig, d *hooks.Dispatcher, deps HookDeps) { + if d == nil { + return + } + tec := &cfg.Batch.ToolExecutorConfig + tec.BeforeToolCall = chainBeforeToolCall(tec.BeforeToolCall, preToolCallHook(d, deps)) + tec.AfterToolCall = chainAfterToolCall(tec.AfterToolCall, postToolCallHook(d, deps)) + installStopHook(cfg, d, deps, "Stop") +} + +// preToolCallHook adapts the dispatcher's PreToolUse event to the BeforeToolCall +// seam. It dispatches with the tool name and raw arguments; a block becomes a +// blocking decision whose reason is surfaced as the tool's error result, and an +// updatedInput becomes an argument rewrite (re-validated by the executor). +func preToolCallHook(d *hooks.Dispatcher, deps HookDeps) agentcore.BeforeToolCallFunc { + return func(ctx context.Context, call agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { + dec := d.Dispatch(ctx, hooks.EventPreToolUse, call.Name, hooks.HookInput{ + EventType: hooks.EventPreToolUse, + SessionID: deps.SessionID, + ProjectDir: deps.ProjectDir, + ToolName: call.Name, + ToolInput: call.Arguments, + }) + if dec.Block { + content := agentcore.ContentList{agentcore.NewTextContent(hookReason(dec.Reason, call.Name))} + return &agentcore.BeforeToolCallDecision{Block: true, Content: &content} + } + if len(dec.UpdatedInput) > 0 { + return &agentcore.BeforeToolCallDecision{UpdatedInput: dec.UpdatedInput} + } + return nil + } +} + +// postToolCallHook adapts the dispatcher's PostToolUse event to the AfterToolCall +// seam. It dispatches with the tool name, raw arguments, and the tool's response, +// then appends any reason/additionalContext to the result content as a new text +// block (the executed tool is never undone). A block on Post is treated as +// feedback only: it cannot retract an already-run tool, so we surface the reason. +func postToolCallHook(d *hooks.Dispatcher, deps HookDeps) agentcore.AfterToolCallFunc { + return func(ctx context.Context, call agentcore.AgentToolCall, result agentcore.AgentToolResult, isError bool) *agentcore.AfterToolCallResult { + var resp json.RawMessage + if b, err := json.Marshal(result.Content); err == nil { + resp = b + } + dec := d.Dispatch(ctx, "PostToolUse", call.Name, hooks.HookInput{ + EventType: "PostToolUse", + SessionID: deps.SessionID, + ProjectDir: deps.ProjectDir, + ToolName: call.Name, + ToolInput: call.Arguments, + ToolResponse: resp, + }) + feedback := joinHookText(dec.Reason, dec.AdditionalContext) + if feedback == "" { + return nil + } + content := append(agentcore.ContentList{}, result.Content...) + content = append(content, agentcore.NewTextContent(feedback)) + return &agentcore.AfterToolCallResult{Content: &content} + } +} + +// hookReason returns the block reason, falling back to a generic message keyed on +// the tool name when the hook gave no reason. +func hookReason(reason, toolName string) string { + if reason != "" { + return reason + } + return "tool " + toolName + " blocked by PreToolUse hook" +} + +// joinHookText joins two hook text fields with a newline, dropping empties. +func joinHookText(a, b string) string { + switch { + case a == "": + return b + case b == "": + return a + default: + return a + "\n" + b + } +} + +// chainBeforeToolCall composes two BeforeToolCall seams into one that runs prev +// first, then next. A blocking decision from prev short-circuits (next does not +// run), so an earlier gate (e.g. trust) is authoritative over a later hook. A +// nil operand is treated as identity, so composing onto an unset seam returns +// the other unchanged. +func chainBeforeToolCall(prev, next agentcore.BeforeToolCallFunc) agentcore.BeforeToolCallFunc { + if prev == nil { + return next + } + if next == nil { + return prev + } + return func(ctx context.Context, call agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { + if dec := prev(ctx, call); dec != nil && dec.Block { + return dec + } + return next(ctx, call) + } +} + +// chainAfterToolCall composes two AfterToolCall seams into one that runs prev +// first, then next. next's non-nil result wins (last writer), so a hook layered +// after an existing seam can override it; when next returns nil, prev's result +// is preserved. A nil operand is identity. +func chainAfterToolCall(prev, next agentcore.AfterToolCallFunc) agentcore.AfterToolCallFunc { + if prev == nil { + return next + } + if next == nil { + return prev + } + return func(ctx context.Context, call agentcore.AgentToolCall, result agentcore.AgentToolResult, isError bool) *agentcore.AfterToolCallResult { + prevRes := prev(ctx, call, result, isError) + if nextRes := next(ctx, call, result, isError); nextRes != nil { + return nextRes + } + return prevRes + } +} + +// chainShouldStop composes two ShouldStopAfterTurn seams with OR semantics: the +// run stops after a turn if either predicate says so. prev runs first and +// short-circuits when true. A nil operand is identity. +func chainShouldStop(prev, next func(context.Context, *agentcore.AgentContext) bool) func(context.Context, *agentcore.AgentContext) bool { + if prev == nil { + return next + } + if next == nil { + return prev + } + return func(ctx context.Context, agentCtx *agentcore.AgentContext) bool { + if prev(ctx, agentCtx) { + return true + } + return next(ctx, agentCtx) + } +} diff --git a/pigo/internal/cli/run/hooks_install_test.go b/pigo/internal/cli/run/hooks_install_test.go new file mode 100644 index 0000000..7069784 --- /dev/null +++ b/pigo/internal/cli/run/hooks_install_test.go @@ -0,0 +1,108 @@ +package run + +import ( + "context" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +func TestInstallHooksEmptyShortCircuits(t *testing.T) { + var cfg runtime.RunConfig + if d := InstallHooks(&cfg, nil, HookDeps{ProjectDir: t.TempDir()}); d != nil { + t.Fatalf("expected nil dispatcher for empty hook set, got %v", d) + } + if d := InstallHooks(&cfg, hooks.HookSet{}, HookDeps{}); d != nil { + t.Fatalf("expected nil dispatcher for empty map, got %v", d) + } +} + +func TestInstallHooksBuildsDispatcher(t *testing.T) { + set := hooks.HookSet{ + "PreToolUse": {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: "true"}}}}, + } + var cfg runtime.RunConfig + d := InstallHooks(&cfg, set, HookDeps{ProjectDir: t.TempDir()}) + if d == nil { + t.Fatal("expected non-nil dispatcher for non-empty hook set") + } +} + +func TestChainBeforeToolCall(t *testing.T) { + block := &agentcore.BeforeToolCallDecision{Block: true} + allow := func(context.Context, agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { return nil } + deny := func(context.Context, agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { return block } + + // nil operands act as identity. + if got := chainBeforeToolCall(nil, deny); got == nil { + t.Fatal("nil prev should return next") + } + if got := chainBeforeToolCall(allow, nil); got == nil { + t.Fatal("nil next should return prev") + } + + // prev blocks → short-circuit, next never runs. + nextRan := false + spy := func(context.Context, agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { + nextRan = true + return nil + } + if dec := chainBeforeToolCall(deny, spy)(context.Background(), agentcore.AgentToolCall{}); dec == nil || !dec.Block { + t.Fatalf("expected block decision, got %v", dec) + } + if nextRan { + t.Fatal("next should not run after prev blocks") + } + + // prev allows → next runs and decides. + if dec := chainBeforeToolCall(allow, deny)(context.Background(), agentcore.AgentToolCall{}); dec == nil || !dec.Block { + t.Fatalf("expected next's block decision, got %v", dec) + } +} + +func TestChainAfterToolCall(t *testing.T) { + content := agentcore.ContentList{} + prevRes := &agentcore.AfterToolCallResult{Content: &content} + nextRes := &agentcore.AfterToolCallResult{Content: &content} + prev := func(context.Context, agentcore.AgentToolCall, agentcore.AgentToolResult, bool) *agentcore.AfterToolCallResult { + return prevRes + } + nilNext := func(context.Context, agentcore.AgentToolCall, agentcore.AgentToolResult, bool) *agentcore.AfterToolCallResult { + return nil + } + next := func(context.Context, agentcore.AgentToolCall, agentcore.AgentToolResult, bool) *agentcore.AfterToolCallResult { + return nextRes + } + + // next returns nil → prev's result preserved. + if got := chainAfterToolCall(prev, nilNext)(context.Background(), agentcore.AgentToolCall{}, agentcore.AgentToolResult{}, false); got != prevRes { + t.Fatalf("expected prev result when next is nil, got %v", got) + } + // next returns non-nil → next wins. + if got := chainAfterToolCall(prev, next)(context.Background(), agentcore.AgentToolCall{}, agentcore.AgentToolResult{}, false); got != nextRes { + t.Fatalf("expected next result to win, got %v", got) + } +} + +func TestChainShouldStop(t *testing.T) { + yes := func(context.Context, *agentcore.AgentContext) bool { return true } + no := func(context.Context, *agentcore.AgentContext) bool { return false } + + if got := chainShouldStop(no, no)(context.Background(), nil); got { + t.Fatal("both false should be false") + } + if got := chainShouldStop(no, yes)(context.Background(), nil); !got { + t.Fatal("next true should stop") + } + // prev true short-circuits without consulting next. + nextRan := false + spy := func(context.Context, *agentcore.AgentContext) bool { nextRan = true; return false } + if got := chainShouldStop(yes, spy)(context.Background(), nil); !got { + t.Fatal("prev true should stop") + } + if nextRan { + t.Fatal("next should not run after prev returns true") + } +} diff --git a/pigo/internal/cli/run/hooks_prompt.go b/pigo/internal/cli/run/hooks_prompt.go new file mode 100644 index 0000000..5f2d703 --- /dev/null +++ b/pigo/internal/cli/run/hooks_prompt.go @@ -0,0 +1,55 @@ +// This file provides the cli-layer helper that runs the UserPromptSubmit hook +// event (US-007, FR-9) at the two prompt entry points (REPL + headless) before a +// prompt is submitted to the loop. It lives in the cli layer alongside the other +// hook assembly (hooks_install.go) because it bridges the resolved Dispatcher and +// the runtime.RunConfig, which runtime itself must not know about. +// +// Two hook effects are supported, with block taking priority over injection: +// +// - block (decision=block / exit 2): the prompt is NOT submitted. The caller +// decides how to surface the reason — the REPL returns to its input state and +// shows it, the headless driver exits non-zero. DispatchUserPromptSubmit only +// reports (block, reason); it does not itself abort. +// - additionalContext: registered as a ONE-SHOT reminder on cfg.Reminders so the +// text is injected into THIS turn's provider context via the existing +// TransformContext seam, then never again. It is not written to the persisted +// message history (the reminder mechanism is ephemeral by construction). +package run + +import ( + "context" + + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// DispatchUserPromptSubmit runs the UserPromptSubmit event for a prompt about to +// be submitted. It returns (block, reason): when block is true the caller must +// NOT submit the prompt and should surface reason. When block is false and the +// hook returned additionalContext, that context is registered as a one-shot +// reminder on cfg.Reminders (allocating a registry when cfg.Reminders is nil) so +// it is injected into this turn only. Block takes priority: a blocking decision +// never also injects. +// +// A nil dispatcher (no hooks configured) is a no-op that returns (false, ""). +func DispatchUserPromptSubmit(ctx context.Context, d *hooks.Dispatcher, cfg *runtime.RunConfig, deps HookDeps, prompt string) (block bool, reason string) { + if d == nil { + return false, "" + } + dec := d.Dispatch(ctx, "UserPromptSubmit", "", hooks.HookInput{ + EventType: "UserPromptSubmit", + SessionID: deps.SessionID, + ProjectDir: deps.ProjectDir, + Prompt: prompt, + }) + if dec.Block { + return true, hookReason(dec.Reason, "UserPromptSubmit") + } + if dec.AdditionalContext != "" && cfg != nil { + if cfg.Reminders == nil { + cfg.Reminders = runtime.NewReminderRegistry() + } + cfg.Reminders.Register(runtime.NewOneShotReminder("user-prompt-submit", dec.AdditionalContext)) + } + return false, "" +} diff --git a/pigo/internal/cli/run/hooks_prompt_test.go b/pigo/internal/cli/run/hooks_prompt_test.go new file mode 100644 index 0000000..e3ec2d6 --- /dev/null +++ b/pigo/internal/cli/run/hooks_prompt_test.go @@ -0,0 +1,97 @@ +package run + +import ( + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// promptDispatcher builds a Dispatcher for a single UserPromptSubmit matcher. +func promptDispatcher(t *testing.T, cmd string) *hooks.Dispatcher { + t.Helper() + set := hooks.HookSet{ + "UserPromptSubmit": {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + d := hooks.NewDispatcher(set, t.TempDir(), nil) + if d == nil { + t.Fatal("expected non-nil dispatcher") + } + return d +} + +// TestUserPromptSubmitBlocks: a hook exiting 2 with a stderr reason blocks the +// prompt; DispatchUserPromptSubmit reports (true, reason) and injects nothing. +func TestUserPromptSubmitBlocks(t *testing.T) { + d := promptDispatcher(t, `echo "prompt rejected" 1>&2; exit 2`) + var cfg runtime.RunConfig + block, reason := DispatchUserPromptSubmit(context.Background(), d, &cfg, HookDeps{SessionID: "s1"}, "hello") + + if !block { + t.Fatal("expected block") + } + if !strings.Contains(reason, "prompt rejected") { + t.Fatalf("block reason not surfaced, got %q", reason) + } + if cfg.Reminders != nil && !cfg.Reminders.Empty() { + t.Fatal("a blocking decision must not register a reminder") + } +} + +// TestUserPromptSubmitInjectsOnce: a hook printing additionalContext registers a +// one-shot reminder that fires exactly once, then goes silent. +func TestUserPromptSubmitInjectsOnce(t *testing.T) { + d := promptDispatcher(t, `echo '{"additionalContext":"remember: run gofmt"}'`) + var cfg runtime.RunConfig + block, _ := DispatchUserPromptSubmit(context.Background(), d, &cfg, HookDeps{SessionID: "s1"}, "hello") + + if block { + t.Fatal("additionalContext must not block") + } + if cfg.Reminders == nil || cfg.Reminders.Empty() { + t.Fatal("additionalContext should register a reminder") + } + + // First turn: the injected context appears. + first := cfg.Reminders.Messages(context.Background(), nil) + if len(first) != 1 { + t.Fatalf("expected 1 reminder message on first turn, got %d", len(first)) + } + if txt := textOfUserMsg(first[0]); !strings.Contains(txt, "remember: run gofmt") { + t.Fatalf("injected context missing, got %q", txt) + } + + // Second turn: the one-shot provider is silent. + if second := cfg.Reminders.Messages(context.Background(), nil); len(second) != 0 { + t.Fatalf("one-shot reminder must not fire twice, got %d", len(second)) + } +} + +// TestUserPromptSubmitNilDispatcher: no hooks configured is a no-op. +func TestUserPromptSubmitNilDispatcher(t *testing.T) { + var cfg runtime.RunConfig + block, reason := DispatchUserPromptSubmit(context.Background(), nil, &cfg, HookDeps{}, "hello") + if block || reason != "" { + t.Fatalf("nil dispatcher should be a no-op, got (%v, %q)", block, reason) + } + if cfg.Reminders != nil { + t.Fatal("nil dispatcher must not allocate a registry") + } +} + +func textOfUserMsg(m agentcore.AgentMessage) string { + um, ok := m.(agentcore.UserMessage) + if !ok { + return "" + } + var b strings.Builder + for _, c := range um.Content { + if tc, ok := c.(agentcore.TextContent); ok { + b.WriteString(tc.Text) + } + } + return b.String() +} diff --git a/pigo/internal/cli/run/hooks_session.go b/pigo/internal/cli/run/hooks_session.go new file mode 100644 index 0000000..da3611c --- /dev/null +++ b/pigo/internal/cli/run/hooks_session.go @@ -0,0 +1,51 @@ +// This file provides the cli-layer helper that runs the SessionStart hook event +// (US-010, FR-1) synchronously at the run-start seam. It lives in the cli layer +// alongside the other hook assembly (hooks_install.go) because it bridges the +// resolved Dispatcher and the runtime.RunConfig, which runtime itself must not +// know about. +// +// SessionStart is dispatched SYNCHRONOUSLY at run start rather than through the +// async OnEvent notifier: an async dispatch could land after the first turn's +// request is built, so its additionalContext would miss the first turn (SPEC +// §11.2). Dispatching inline before the loop starts guarantees the injected +// context is present for the very first turn. +// +// Only injection is supported (there is nothing to block at session start): any +// additionalContext is registered as a ONE-SHOT reminder on cfg.Reminders so the +// text is injected into the first turn's provider context via the existing +// TransformContext seam, then never again. It is not written to the persisted +// message history (the reminder mechanism is ephemeral by construction). +package run + +import ( + "context" + + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// DispatchSessionStart runs the SessionStart event once at run start. source is +// "startup" for a fresh run or "resume" when continuing an existing session; it +// is carried in the HookInput so hooks can differentiate. Any additionalContext +// returned by the hook is registered as a one-shot reminder on cfg.Reminders +// (allocating a registry when cfg.Reminders is nil) so it is injected into the +// first turn only. +// +// A nil dispatcher (no hooks configured) is a no-op. +func DispatchSessionStart(ctx context.Context, d *hooks.Dispatcher, cfg *runtime.RunConfig, deps HookDeps, source string) { + if d == nil { + return + } + dec := d.Dispatch(ctx, "SessionStart", "", hooks.HookInput{ + EventType: "SessionStart", + SessionID: deps.SessionID, + ProjectDir: deps.ProjectDir, + Source: source, + }) + if dec.AdditionalContext != "" && cfg != nil { + if cfg.Reminders == nil { + cfg.Reminders = runtime.NewReminderRegistry() + } + cfg.Reminders.Register(runtime.NewOneShotReminder("session-start", dec.AdditionalContext)) + } +} diff --git a/pigo/internal/cli/run/hooks_session_test.go b/pigo/internal/cli/run/hooks_session_test.go new file mode 100644 index 0000000..7748658 --- /dev/null +++ b/pigo/internal/cli/run/hooks_session_test.go @@ -0,0 +1,76 @@ +package run + +import ( + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// sessionDispatcher builds a Dispatcher for a single SessionStart matcher. +func sessionDispatcher(t *testing.T, cmd string) *hooks.Dispatcher { + t.Helper() + set := hooks.HookSet{ + "SessionStart": {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + d := hooks.NewDispatcher(set, t.TempDir(), nil) + if d == nil { + t.Fatal("expected non-nil dispatcher") + } + return d +} + +// TestSessionStartInjectsOnce: a SessionStart hook printing additionalContext +// registers a one-shot reminder that fires on the first turn, then goes silent. +func TestSessionStartInjectsOnce(t *testing.T) { + d := sessionDispatcher(t, `echo '{"additionalContext":"project context loaded"}'`) + var cfg runtime.RunConfig + DispatchSessionStart(context.Background(), d, &cfg, HookDeps{SessionID: "s1"}, "startup") + + if cfg.Reminders == nil || cfg.Reminders.Empty() { + t.Fatal("SessionStart additionalContext should register a reminder") + } + + first := cfg.Reminders.Messages(context.Background(), nil) + if len(first) != 1 { + t.Fatalf("expected 1 reminder message on the first turn, got %d", len(first)) + } + if txt := textOfUserMsg(first[0]); !strings.Contains(txt, "project context loaded") { + t.Fatalf("injected context missing, got %q", txt) + } + + if second := cfg.Reminders.Messages(context.Background(), nil); len(second) != 0 { + t.Fatalf("one-shot reminder must not fire twice, got %d", len(second)) + } +} + +// TestSessionStartResumeSource: the source ("resume") is threaded into the hook +// input so the hook can differentiate startup from resume. The hook echoes its +// stdin JSON's source field into additionalContext, which we then observe. +func TestSessionStartResumeSource(t *testing.T) { + d := sessionDispatcher(t, `in=$(cat); case "$in" in *'"source":"resume"'*) echo '{"additionalContext":"source=resume"}';; *) echo '{}';; esac`) + var cfg runtime.RunConfig + DispatchSessionStart(context.Background(), d, &cfg, HookDeps{SessionID: "s1"}, "resume") + + if cfg.Reminders == nil || cfg.Reminders.Empty() { + t.Fatal("expected a reminder registered") + } + msgs := cfg.Reminders.Messages(context.Background(), nil) + if len(msgs) != 1 { + t.Fatalf("expected 1 reminder message, got %d", len(msgs)) + } + if txt := textOfUserMsg(msgs[0]); !strings.Contains(txt, "source=resume") { + t.Fatalf("resume source not threaded into hook input, got %q", txt) + } +} + +// TestSessionStartNilDispatcher: no hooks configured is a no-op. +func TestSessionStartNilDispatcher(t *testing.T) { + var cfg runtime.RunConfig + DispatchSessionStart(context.Background(), nil, &cfg, HookDeps{}, "startup") + if cfg.Reminders != nil { + t.Fatal("nil dispatcher must not allocate a registry") + } +} diff --git a/pigo/internal/cli/run/hooks_stop.go b/pigo/internal/cli/run/hooks_stop.go new file mode 100644 index 0000000..af6a8db --- /dev/null +++ b/pigo/internal/cli/run/hooks_stop.go @@ -0,0 +1,105 @@ +// This file provides the cli-layer decorator that runs the Stop and SubagentStop +// hook events (US-008/009, FR-10/12) at the loop's natural-end seam +// (runtime.RunConfig.OnStop). It lives in the cli layer because it bridges the +// resolved Dispatcher and runtime, which must not depend on hooks. +// +// A Stop hook may block the run from ending and force a continuation, feeding +// its reason back as guidance. Left unchecked a hook that always blocks would +// loop forever, so the decorator holds a per-run consecutive-block counter and +// force-stops after maxConsecutiveStopBlocks (FR-12). The counter resets on any +// natural (non-blocking) stop, so an occasional block does not erode the budget. +package run + +import ( + "context" + "fmt" + "os" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// maxConsecutiveStopBlocks is the default FR-12 ceiling on how many times a Stop +// (or SubagentStop) hook may consecutively block the run from ending before the +// decorator forces a stop and warns. Overridable via the builder. +const maxConsecutiveStopBlocks = 5 + +// stopHook builds a runtime.OnStop seam that dispatches the given event +// ("Stop" for a top-level run, "SubagentStop" inside a sub-agent) each time the +// loop is about to end. On a blocking decision it returns a StopDecision that +// keeps the run alive with the hook's reason as guidance, up to maxBlocks +// consecutive blocks; past that it forces a stop and warns. maxBlocks <= 0 uses +// maxConsecutiveStopBlocks. A nil dispatcher yields a nil seam (no wrapping). +func stopHook(d *hooks.Dispatcher, deps HookDeps, event string, maxBlocks int) func(context.Context, *agentcore.AgentContext) *runtime.StopDecision { + if d == nil { + return nil + } + if maxBlocks <= 0 { + maxBlocks = maxConsecutiveStopBlocks + } + warn := deps.WarnLog + if warn == nil { + warn = os.Stderr + } + consecutive := 0 + return func(ctx context.Context, _ *agentcore.AgentContext) *runtime.StopDecision { + dec := d.Dispatch(ctx, event, "", hooks.HookInput{ + EventType: event, + SessionID: deps.SessionID, + ProjectDir: deps.ProjectDir, + StopReason: "end_turn", + }) + if !dec.Block { + consecutive = 0 + return nil + } + consecutive++ + if consecutive > maxBlocks { + fmt.Fprintf(warn, "%s hook blocked the run %d times consecutively; forcing stop\n", event, consecutive-1) + consecutive = 0 + return nil + } + return &runtime.StopDecision{Block: true, Guidance: hookReason(dec.Reason, event)} + } +} + +// InstallSubagentStop wires the SubagentStop hook onto a sub-agent's RunConfig, +// so a SubagentStop hook can block a sub-agent from ending (same semantics as +// Stop, evaluated in the sub-agent context, with its own consecutive-block +// budget). The sub-agent assembly calls this on the child runCfg (converged in +// #425). A nil dispatcher is a no-op. +func InstallSubagentStop(cfg *runtime.RunConfig, d *hooks.Dispatcher, deps HookDeps) { + installStopHook(cfg, d, deps, "SubagentStop") +} + +// installStopHook wires a Stop-family decorator onto cfg.OnStop, chaining onto +// any existing seam (an earlier OnStop is consulted first and its block wins, +// mirroring the block-short-circuit combinator convention). It is called by the +// tool-execution wiring for the top-level run and by the sub-agent assembly for +// the SubagentStop variant. +func installStopHook(cfg *runtime.RunConfig, d *hooks.Dispatcher, deps HookDeps, event string) { + next := stopHook(d, deps, event, 0) + if next == nil { + return + } + cfg.OnStop = chainOnStop(cfg.OnStop, next) +} + +// chainOnStop composes two OnStop seams: prev is consulted first and a blocking +// decision short-circuits (next does not run), so an earlier gate stays +// authoritative. A nil operand is identity. +func chainOnStop(prev, next func(context.Context, *agentcore.AgentContext) *runtime.StopDecision) func(context.Context, *agentcore.AgentContext) *runtime.StopDecision { + if prev == nil { + return next + } + if next == nil { + return prev + } + return func(ctx context.Context, agentCtx *agentcore.AgentContext) *runtime.StopDecision { + if dec := prev(ctx, agentCtx); dec != nil && dec.Block { + return dec + } + return next(ctx, agentCtx) + } +} diff --git a/pigo/internal/cli/run/hooks_stop_test.go b/pigo/internal/cli/run/hooks_stop_test.go new file mode 100644 index 0000000..1eb864d --- /dev/null +++ b/pigo/internal/cli/run/hooks_stop_test.go @@ -0,0 +1,84 @@ +package run + +import ( + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/hooks" +) + +func stopDispatcher(t *testing.T, event, cmd string) *hooks.Dispatcher { + t.Helper() + set := hooks.HookSet{ + event: {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + d := hooks.NewDispatcher(set, t.TempDir(), nil) + if d == nil { + t.Fatal("expected non-nil dispatcher") + } + return d +} + +// TestStopHookBlocksThenForceStops: a Stop hook that always blocks (exit 2) is +// honored up to the limit, then the decorator force-stops (returns nil) so the +// run cannot loop forever (FR-12). +func TestStopHookBlocksThenForceStops(t *testing.T) { + d := stopDispatcher(t, "Stop", `echo "not done yet" 1>&2; exit 2`) + var warn strings.Builder + seam := stopHook(d, HookDeps{SessionID: "s1", WarnLog: &warn}, "Stop", 3) + if seam == nil { + t.Fatal("expected non-nil seam") + } + + // First 3 consultations block with the reason as guidance. + for i := 0; i < 3; i++ { + dec := seam(context.Background(), nil) + if dec == nil || !dec.Block { + t.Fatalf("consult %d: expected a blocking decision", i+1) + } + if !strings.Contains(dec.Guidance, "not done yet") { + t.Fatalf("consult %d: block reason not surfaced as guidance, got %q", i+1, dec.Guidance) + } + } + // 4th consultation exceeds the limit: force stop (nil) + a warning. + if dec := seam(context.Background(), nil); dec != nil { + t.Fatalf("expected force-stop (nil) past the limit, got %+v", dec) + } + if !strings.Contains(warn.String(), "forcing stop") { + t.Fatalf("force-stop should warn, got %q", warn.String()) + } +} + +// TestStopHookResetsCounterOnAllow: a non-blocking hook always returns nil, so +// the run is free to end and the consecutive-block counter never accrues. +func TestStopHookResetsCounterOnAllow(t *testing.T) { + d := stopDispatcher(t, "Stop", `exit 0`) + seam := stopHook(d, HookDeps{SessionID: "s1"}, "Stop", 2) + for i := 0; i < 5; i++ { + if dec := seam(context.Background(), nil); dec != nil { + t.Fatalf("consult %d: a non-blocking hook must let the run end (nil), got %+v", i+1, dec) + } + } +} + +// TestStopHookNilDispatcher: no hooks configured yields a nil seam (no wrapping). +func TestStopHookNilDispatcher(t *testing.T) { + if seam := stopHook(nil, HookDeps{}, "Stop", 0); seam != nil { + t.Fatal("nil dispatcher must yield a nil seam") + } +} + +// TestSubagentStopEvent: InstallSubagentStop wires a SubagentStop decorator whose +// block keeps a sub-agent running with the hook's reason. +func TestSubagentStopEvent(t *testing.T) { + d := stopDispatcher(t, "SubagentStop", `echo "sub not done" 1>&2; exit 2`) + seam := stopHook(d, HookDeps{SessionID: "child"}, "SubagentStop", 5) + if seam == nil { + t.Fatal("expected non-nil SubagentStop seam") + } + dec := seam(context.Background(), nil) + if dec == nil || !dec.Block || !strings.Contains(dec.Guidance, "sub not done") { + t.Fatalf("SubagentStop block not honored, got %+v", dec) + } +} diff --git a/pigo/internal/cli/run/hooks_tooluse_test.go b/pigo/internal/cli/run/hooks_tooluse_test.go new file mode 100644 index 0000000..fb95609 --- /dev/null +++ b/pigo/internal/cli/run/hooks_tooluse_test.go @@ -0,0 +1,189 @@ +package run + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/runtime" +) + +// recordingTool is a fake AgentTool that records whether Execute ran and echoes +// a fixed result, so a PreToolUse block can be asserted as "never executed". +type recordingTool struct { + name string + ran *bool +} + +func (t recordingTool) Name() string { return t.name } +func (t recordingTool) Description() string { return "fake" } +func (t recordingTool) Schema() json.RawMessage { return nil } +func (t recordingTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} +func (t recordingTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + *t.ran = true + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("executed")}}, nil +} + +func wiredConfig(t *testing.T, tool agentcore.AgentTool, set hooks.HookSet) runtime.RunConfig { + t.Helper() + reg := agenttool.NewToolRegistry() + if err := reg.Register(tool); err != nil { + t.Fatalf("register: %v", err) + } + var cfg runtime.RunConfig + cfg.Batch.ToolExecutorConfig.Registry = reg + if d := InstallHooks(&cfg, set, HookDeps{SessionID: "s1", ProjectDir: t.TempDir()}); d == nil { + t.Fatal("expected non-nil dispatcher") + } + return cfg +} + +// TestPreToolUseBlocksBashRmRf: a PreToolUse hook matching bash inspects the +// piped tool_input for "rm -rf" and exits 2 with a reason; the tool must not run +// and the reason must surface in the result the model receives. +func TestPreToolUseBlocksBashRmRf(t *testing.T) { + ran := false + tool := recordingTool{name: "bash", ran: &ran} + // Hook: block (exit 2) when stdin JSON contains "rm -rf", printing a reason. + cmd := `if grep -q "rm -rf" ; then echo "dangerous command blocked" 1>&2; exit 2; fi` + set := hooks.HookSet{ + "PreToolUse": {{Matcher: "bash", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + cfg := wiredConfig(t, tool, set) + + call := agentcore.AgentToolCall{ID: "1", Name: "bash", Arguments: json.RawMessage(`{"command":"rm -rf /tmp/x"}`)} + msgs, _ := agenttool.ExecuteToolCalls(context.Background(), cfg.Batch, []agentcore.AgentToolCall{call}, nil) + + if ran { + t.Fatal("tool must not execute when PreToolUse blocks") + } + if len(msgs) != 1 || !msgs[0].IsError { + t.Fatalf("blocked call should be an error result: %+v", msgs) + } + if txt := textOfMsg(msgs[0]); !strings.Contains(txt, "dangerous command blocked") { + t.Fatalf("block reason not surfaced to model, got %q", txt) + } +} + +// TestPreToolUseAllowsSafeBash: the same hook allows a command without "rm -rf". +func TestPreToolUseAllowsSafeBash(t *testing.T) { + ran := false + tool := recordingTool{name: "bash", ran: &ran} + cmd := `if grep -q "rm -rf" ; then echo "blocked" 1>&2; exit 2; fi` + set := hooks.HookSet{ + "PreToolUse": {{Matcher: "bash", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + cfg := wiredConfig(t, tool, set) + + call := agentcore.AgentToolCall{ID: "1", Name: "bash", Arguments: json.RawMessage(`{"command":"ls"}`)} + msgs, _ := agenttool.ExecuteToolCalls(context.Background(), cfg.Batch, []agentcore.AgentToolCall{call}, nil) + + if !ran { + t.Fatal("safe command should execute") + } + if msgs[0].IsError { + t.Fatalf("safe command should not error: %+v", msgs[0]) + } +} + +// TestPostToolUseAppendsFeedback: a PostToolUse hook prints additionalContext, +// which must be appended to the executed tool's result (not undo it). +func TestPostToolUseAppendsFeedback(t *testing.T) { + ran := false + tool := recordingTool{name: "write", ran: &ran} + cmd := `echo '{"additionalContext":"linted: 0 issues"}'` + set := hooks.HookSet{ + "PostToolUse": {{Matcher: "write", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + cfg := wiredConfig(t, tool, set) + + call := agentcore.AgentToolCall{ID: "1", Name: "write", Arguments: json.RawMessage(`{"path":"a.go"}`)} + msgs, _ := agenttool.ExecuteToolCalls(context.Background(), cfg.Batch, []agentcore.AgentToolCall{call}, nil) + + if !ran { + t.Fatal("tool should execute; Post hook must not undo it") + } + txt := allTextOfMsg(msgs[0]) + if !strings.Contains(txt, "executed") { + t.Fatalf("original result lost: %q", txt) + } + if !strings.Contains(txt, "linted: 0 issues") { + t.Fatalf("Post hook feedback not appended: %q", txt) + } +} + +// TestPreToolUseUpdatedInputRewritesArgs: a PreToolUse hook returns updatedInput, +// which must replace the tool's arguments before execution. +func TestPreToolUseUpdatedInputRewritesArgs(t *testing.T) { + var gotArgs json.RawMessage + captured := false + tool := capturingTool{name: "bash", got: &gotArgs, captured: &captured} + reg := agenttool.NewToolRegistry() + if err := reg.Register(tool); err != nil { + t.Fatalf("register: %v", err) + } + cmd := `echo '{"updatedInput":{"command":"echo safe"}}'` + set := hooks.HookSet{ + "PreToolUse": {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: cmd}}}}, + } + var cfg runtime.RunConfig + cfg.Batch.ToolExecutorConfig.Registry = reg + if d := InstallHooks(&cfg, set, HookDeps{ProjectDir: t.TempDir()}); d == nil { + t.Fatal("expected non-nil dispatcher") + } + + call := agentcore.AgentToolCall{ID: "1", Name: "bash", Arguments: json.RawMessage(`{"command":"rm -rf /"}`)} + agenttool.ExecuteToolCalls(context.Background(), cfg.Batch, []agentcore.AgentToolCall{call}, nil) + + if !captured { + t.Fatal("tool should have executed with rewritten args") + } + if !strings.Contains(string(gotArgs), "echo safe") { + t.Fatalf("args not rewritten by updatedInput, got %q", string(gotArgs)) + } +} + +type capturingTool struct { + name string + got *json.RawMessage + captured *bool +} + +func (t capturingTool) Name() string { return t.name } +func (t capturingTool) Description() string { return "capture" } +func (t capturingTool) Schema() json.RawMessage { return nil } +func (t capturingTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} +func (t capturingTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + *t.got = args + *t.captured = true + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, nil +} + +func textOfMsg(msg agentcore.ToolResultMessage) string { + if len(msg.Content) == 0 { + return "" + } + if tc, ok := msg.Content[0].(agentcore.TextContent); ok { + return tc.Text + } + return "" +} + +func allTextOfMsg(msg agentcore.ToolResultMessage) string { + var b strings.Builder + for _, c := range msg.Content { + if tc, ok := c.(agentcore.TextContent); ok { + b.WriteString(tc.Text) + b.WriteString("\n") + } + } + return b.String() +} diff --git a/pigo/internal/cli/run/memory_wiring_test.go b/pigo/internal/cli/run/memory_wiring_test.go new file mode 100644 index 0000000..7d957ac --- /dev/null +++ b/pigo/internal/cli/run/memory_wiring_test.go @@ -0,0 +1,87 @@ +package run + +// Tests for the persistent-memory loop wiring (#481): OpenMemoryStore's +// enabled/disabled contract, MemoryRootFromTools resolving through the opened +// store, and TodoReminders registering the memory reminder provider alongside +// the todo one. + +import ( + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/memory" +) + +// TestOpenMemoryStoreDisabled verifies memory.enabled=false yields (nil, nil): +// a disabled store is not an error, so the caller degrades to file-based +// auto-memory without logging a failure. +func TestOpenMemoryStoreDisabled(t *testing.T) { + store, err := OpenMemoryStore(false) + if err != nil { + t.Fatalf("OpenMemoryStore(false) err = %v, want nil", err) + } + if store != nil { + t.Fatalf("OpenMemoryStore(false) store = %v, want nil", store) + } +} + +// TestMemoryDirHonorsPIGOHome verifies MemoryDir roots the store at +// $PIGO_HOME/memory when the override is set. +func TestMemoryDirHonorsPIGOHome(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_HOME", dir) + if got, want := MemoryDir(), filepath.Join(dir, "memory"); got != want { + t.Errorf("MemoryDir() = %q, want %q", got, want) + } +} + +// TestMemoryRootFromToolsPresent verifies the root is resolved through the +// memory_search tool's Store.Root() when one is wired into the tool set. +func TestMemoryRootFromToolsPresent(t *testing.T) { + root := t.TempDir() + store, err := memory.Open(filepath.Join(root, "index.db"), root, "") + if err != nil { + t.Fatalf("memory.Open: %v", err) + } + defer store.Close() + + tools := []agentcore.AgentTool{&agenttool.MemorySearchTool{Store: store}} + if got := MemoryRootFromTools(tools); got != root { + t.Errorf("MemoryRootFromTools = %q, want %q", got, root) + } +} + +// TestMemoryRootFromToolsAbsent verifies the root is "" when no memory_search +// tool (or a store-less one) is present. +func TestMemoryRootFromToolsAbsent(t *testing.T) { + if got := MemoryRootFromTools(nil); got != "" { + t.Errorf("MemoryRootFromTools(nil) = %q, want empty", got) + } + tools := []agentcore.AgentTool{&agenttool.MemorySearchTool{Store: nil}} + if got := MemoryRootFromTools(tools); got != "" { + t.Errorf("MemoryRootFromTools(store-less) = %q, want empty", got) + } +} + +// TestTodoRemindersRegistersMemoryProvider verifies TodoReminders builds a +// non-empty registry from a memory_search tool alone, and stays nil when no +// provider-bearing tool is present. +func TestTodoRemindersRegistersMemoryProvider(t *testing.T) { + root := t.TempDir() + store, err := memory.Open(filepath.Join(root, "index.db"), root, "") + if err != nil { + t.Fatalf("memory.Open: %v", err) + } + defer store.Close() + + reg := TodoReminders([]agentcore.AgentTool{&agenttool.MemorySearchTool{Store: store}}) + if reg == nil || reg.Empty() { + t.Fatal("TodoReminders with a memory_search tool should yield a non-empty registry") + } + + if reg := TodoReminders(nil); reg != nil { + t.Errorf("TodoReminders(nil) = %v, want nil", reg) + } +} diff --git a/pigo/internal/cli/run/prompt_flags_test.go b/pigo/internal/cli/run/prompt_flags_test.go new file mode 100644 index 0000000..f7370c9 --- /dev/null +++ b/pigo/internal/cli/run/prompt_flags_test.go @@ -0,0 +1,88 @@ +package run + +// Tests for --append-system-prompt value resolution (mirrors pi): each value is +// either a path to an existing file whose contents are appended, or literal +// text when it is not an existing file. A value that names an unreadable file +// (a real I/O error other than not-exist) is surfaced rather than silently +// appended verbatim. + +import ( + "os" + "path/filepath" + "testing" +) + +// TestResolveAppendInstructionsEmpty verifies no values yields no appends. +func TestResolveAppendInstructionsEmpty(t *testing.T) { + out, err := resolveAppendInstructions(nil) + if err != nil { + t.Fatalf("resolveAppendInstructions: %v", err) + } + if out != nil { + t.Errorf("expected nil for no values, got %#v", out) + } +} + +// TestResolveAppendInstructionsLiteral verifies a value that is not an existing +// file is treated as literal text and passed through verbatim. +func TestResolveAppendInstructionsLiteral(t *testing.T) { + out, err := resolveAppendInstructions([]string{"be concise and helpful"}) + if err != nil { + t.Fatalf("resolveAppendInstructions: %v", err) + } + if len(out) != 1 || out[0] != "be concise and helpful" { + t.Errorf("literal text must pass through verbatim, got %#v", out) + } +} + +// TestResolveAppendInstructionsFile verifies a value that names an existing +// file has its contents read and appended. +func TestResolveAppendInstructionsFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "guidance.txt") + if err := os.WriteFile(path, []byte("FILE GUIDANCE"), 0o644); err != nil { + t.Fatalf("write temp file: %v", err) + } + out, err := resolveAppendInstructions([]string{path, "literal tail"}) + if err != nil { + t.Fatalf("resolveAppendInstructions: %v", err) + } + if len(out) != 2 { + t.Fatalf("expected 2 resolved values, got %#v", out) + } + if out[0] != "FILE GUIDANCE" { + t.Errorf("existing file must be read into the append, got %q", out[0]) + } + if out[1] != "literal tail" { + t.Errorf("literal value must pass through verbatim, got %q", out[1]) + } +} + +// TestResolveAppendInstructionsUnreadableFile verifies a value that looks like a +// path but points at an unreadable file (here, a directory) surfaces an error +// rather than being appended verbatim. A directory is used because os.Stat +// succeeds on it (so it is not treated as literal text) while os.ReadFile fails. +func TestResolveAppendInstructionsUnreadableFile(t *testing.T) { + dir := t.TempDir() + // A directory: Stat succeeds and IsDir() is true, so it is treated as literal + // text — assert that. Then create an actually-unreadable regular file to hit + // the read-error path. + if out, err := resolveAppendInstructions([]string{dir}); err != nil { + t.Fatalf("a directory should be treated as literal text, got error: %v", err) + } else if len(out) != 1 || out[0] != dir { + t.Errorf("a directory path must pass through as literal text, got %#v", out) + } + + unreadable := filepath.Join(dir, "secret.txt") + if err := os.WriteFile(unreadable, []byte("nope"), 0o000); err != nil { + t.Fatalf("write temp file: %v", err) + } + // Root can read 0o000 files, so skip the read-error assertion when running as + // root (common in CI containers) — the path would succeed instead of erroring. + if os.Geteuid() == 0 { + t.Skip("running as root: 0o000 file is still readable, cannot exercise read-error path") + } + if _, err := resolveAppendInstructions([]string{unreadable}); err == nil { + t.Error("an unreadable append file must surface an error, got nil") + } +} diff --git a/pigo/internal/cli/run/run.go b/pigo/internal/cli/run/run.go new file mode 100644 index 0000000..7e78140 --- /dev/null +++ b/pigo/internal/cli/run/run.go @@ -0,0 +1,604 @@ +// Package run holds the run-assembly layer (US-005, #362): the shared setup that +// both the interactive REPL and the headless driver need — resolving the +// provider, building the tool set rooted at the working directory, discovering +// skills and plugins, and constructing the loop RunConfig. Pulling it out of +// cmd/pigo lets the subpackages assemble a run through one exported API instead +// of duplicating the wiring. +package run + +import ( + "fmt" + "io" + "os" + "path/filepath" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/builtinskills" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/memory" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/trust" +) + +// Env is the environment every run shares: the working directory, the tool set +// rooted at it, the resolved provider, and the system prompt. It is assembled +// once (SetupEnv) and consumed by whichever driver runs. +type Env struct { + Cwd string + Tools []agentcore.AgentTool + Provider provider.Provider + ProviderName string + SysPrompt string + + // Skills is the discovered skill set (loaded once here, empty under + // --no-skills). It is threaded into the REPL so each skill is registered as a + // /skill-name command, and the model-invocable subset is already injected into + // SysPrompt. + Skills []*runtime.Skill + + // Plugins holds any loaded external plugins so the caller can Close them when + // the run ends. It is nil when no plugins were discovered. + Plugins *plugin.Manager + + // Memory is the persistent memory store opened once for the run (issue #481), + // or nil when persistent memory is disabled (memory.enabled=false), tools are + // disabled (--no-tools), or the store could not be opened (a non-fatal + // failure). When non-nil the caller MUST Close it when the run ends. The store + // is also handed to the memory_search tool (in Tools) and, through it, the + // per-turn memory reminder provider, so this field exists mainly so the owner + // can close the DB — downstream wiring reaches the store via the tool. + Memory *memory.Store +} + +// SetupEnv resolves the provider for model/baseURL, builds the tool set rooted +// at the working directory, and constructs the system prompt — the setup the +// REPL and headless drivers both need. systemPrompt, when non-empty, replaces +// the default base instruction (mirrors pi's --system-prompt); appendSystemPrompt +// entries are each resolved (a path to an existing file is read, otherwise the +// value is literal text) and layered onto the end of the prompt (mirrors pi's +// --append-system-prompt). apiKey is the resolved credential (CLI --api-key or +// config.toml) used as the override for sub-agent credential resolution so +// dispatched task children authenticate the same way the parent does. policy is +// the --allowed-tools/--disallowed-tools boundary; it is validated against the +// fully assembled tool set and then applied, so an unknown tool name is a usage +// error rather than a silently ineffective boundary. It returns an error rather +// than exiting so the caller owns exit-code mapping. +func SetupEnv(model, baseURL, protocol, providerName, apiKey string, noTools, noSkills bool, systemPrompt string, appendSystemPrompt []string, memEnabled bool, policy ToolPolicy) (Env, error) { + cwd, _ := os.Getwd() + prov, resolvedName, err := provider.ResolveProvider(model, baseURL, protocol, providerName, os.Getenv) + if err != nil { + return Env{}, err + } + appends, err := resolveAppendInstructions(appendSystemPrompt) + if err != nil { + return Env{}, err + } + tools := BuiltinTools(cwd, noTools) + // Open the persistent memory store once (issue #481) and expose it as the + // memory_search tool so the agent can recall earlier context. Memory is a + // tool, so it is skipped under --no-tools; memory.enabled=false disables it + // too. Opening is non-fatal: a failure logs and leaves memory off, matching + // the "fall back to file-based auto-memory" contract. + var memStore *memory.Store + if !noTools { + if store, err := OpenMemoryStore(memEnabled); err != nil { + fmt.Fprintf(os.Stderr, "pigo: memory disabled: %v\n", err) + } else if store != nil { + memStore = store + tools = append(tools, &agenttool.MemorySearchTool{Store: store}) + } + } + // Wire the generic task tool (US-002, #454) unless tools are disabled. It + // dispatches general-purpose sub-agents that reuse the resolved provider + // stream/model. Each spawn gets a fresh child RunConfig whose registry is the + // builtins with "task" removed (the nesting guard, so a child cannot fan out + // again), and all task calls in a run share one semaphore capping concurrency. + if !noTools { + sem := runtime.NewSubagentSemaphore() + // The child resolves credentials the same way the parent does: env/OAuth via + // a fresh store, plus the CLI/config api key as an override. Without the + // override a child would get an empty key whenever auth comes from config.toml + // or --api-key (not an env var), leaving every sub-agent unauthenticated. + childCreds := provider.NewCredentialStore(nil) + childCreds.SetOverride(resolvedName, apiKey) + factory := func() runtime.RunConfig { + childTools := ChildToolSet(cwd, policy) + return runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: model, + Provider: resolvedName, + Stream: provider.StreamFnFromProvider(prov), + GetAPIKey: childCreds.GetAPIKey, + }, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: ToolRegistry(childTools)}}, + } + } + tools = append(tools, runtime.NewTaskTool(factory, sem)) + } + // Wire the blackboard tool (coop/): present only when the BB environment + // variable names a blackboard root. It is the atomic shared-file primitive + // of the pigo coop runner (task.md / workspace / DONE); without BB it is + // absent so ordinary runs are unaffected. Like memory, it is a tool, so + // --no-tools disables it. + if !noTools { + if bb := strings.TrimSpace(os.Getenv("BB")); bb != "" { + tools = append(tools, &agenttool.BlackboardTool{Root: bb}) + } + } + // Discover external plugins (US-016) and append their tools. Plugin loading + // is fault-tolerant: a plugin that fails to start is logged and skipped, and + // disabling tools (--no-tools) skips plugin discovery entirely. + var mgr *plugin.Manager + if !noTools { + if m, err := plugin.Discover(PluginsDir(), os.Stderr, os.Stderr); err == nil { + tools = append(tools, m.Tools()...) + mgr = m + } else { + fmt.Fprintf(os.Stderr, "pigo: plugin discovery failed: %v\n", err) + } + } + // Enforce the --allowed-tools/--disallowed-tools boundary now that the set is + // complete. Validation must happen here rather than at flag-parse time: plugin + // and memory tool names only exist at runtime, so an earlier check would reject + // legitimate names. Filtering here — at the registration layer, before the + // BeforeToolCall confirmation gate — is what makes the boundary structural: a + // removed tool is never advertised and never dispatchable, so --approve cannot + // widen it. + if err := ValidateToolPolicy(tools, policy); err != nil { + return Env{}, err + } + tools = ApplyToolPolicy(tools, policy) + if len(tools) == 0 && !noTools && !policy.IsZero() { + fmt.Fprintln(os.Stderr, "pigo: warning: the tool policy removed every tool; the model will run without tools") + } + // --no-tools already disables everything, so a tool policy alongside it has + // no effect — and because the set is empty, ValidateToolPolicy above skipped + // name validation, meaning a typo here would otherwise pass unnoticed. Say so + // rather than letting the user believe a boundary is in force. + if noTools && !policy.IsZero() { + fmt.Fprintln(os.Stderr, "pigo: warning: --no-tools disables all tools; --allowed-tools/--disallowed-tools are ignored (and unvalidated)") + } + // Load skills once (shared between prompt injection and /skill-name + // registration). A partial parse error still yields the skills that DID load, + // so one malformed file is a non-fatal warning rather than a hard failure. + skills, err := LoadSkills(noSkills) + if err != nil { + fmt.Fprintf(os.Stderr, "pigo: skills: %v\n", err) + } + // The model can only load a skill's body when the read tool is present, so + // advertise skills in the prompt only then (mirrors pi's selectedTools check). + sysPrompt, err := runtime.BuildSystemPrompt(runtime.PromptConfig{ + BaseInstruction: systemPrompt, + WorkingDir: cwd, + Root: cwd, + AppendInstructions: appends, + Skills: skills, + ReadToolAvailable: hasReadTool(tools), + }) + if err != nil { + return Env{}, err + } + return Env{ + Cwd: cwd, + Tools: tools, + Provider: prov, + ProviderName: resolvedName, + SysPrompt: sysPrompt, + Skills: skills, + Plugins: mgr, + Memory: memStore, + }, nil +} + +// hasReadTool reports whether the read tool is present in the tool set. Skills +// are advertised in the system prompt only when it is, since the model needs the +// read tool to load a skill's body on demand. +func hasReadTool(tools []agentcore.AgentTool) bool { + for _, t := range tools { + if t.Name() == "read" { + return true + } + } + return false +} + +// resolveAppendInstructions maps each --append-system-prompt value to the text +// to append. Following pi, a value that names an existing regular file is read +// and its contents are appended; any other value (a non-existent path, or a +// directory) is treated as literal text. Only a value that stats as a regular +// file but then fails to read (e.g. a permission error) is reported, so a +// genuinely broken file path is not silently appended verbatim. +func resolveAppendInstructions(values []string) ([]string, error) { + if len(values) == 0 { + return nil, nil + } + out := make([]string, 0, len(values)) + for _, v := range values { + info, statErr := os.Stat(v) + if statErr == nil && !info.IsDir() { + data, err := os.ReadFile(v) + if err != nil { + return nil, fmt.Errorf("read --append-system-prompt file %q: %w", v, err) + } + out = append(out, string(data)) + continue + } + out = append(out, v) + } + return out, nil +} + +// BuiltinTools returns the default file/shell tool set rooted at cwd, or nil +// when tools are disabled. The todo tool is stateful: a single TodoStore is +// created here and held by the one TodoTool instance, so the task list persists +// across calls within a run (a later write replaces the plan). +func BuiltinTools(cwd string, disabled bool) []agentcore.AgentTool { + if disabled { + return nil + } + // A single recorder is shared by the write and edit tools so /rewind can roll + // back every mutation from a turn regardless of which tool made it. + snap := agenttool.NewFileSnapshotRecorder() + // A single job store is shared by bash, bash_output and kill_bash so a + // background command launched by bash is visible to the drain/kill tools. + jobs := agenttool.NewBashJobStore() + return []agentcore.AgentTool{ + &agenttool.ReadTool{Root: cwd, ExtraRoots: ReadableExtraRoots()}, + &agenttool.WriteTool{Root: cwd, ExtraRoots: ReadableExtraRoots(), Snap: snap}, + &agenttool.EditTool{Root: cwd, ExtraRoots: ReadableExtraRoots(), Snap: snap}, + &agenttool.GrepTool{Root: cwd}, + &agenttool.FindTool{Root: cwd}, + &agenttool.BashTool{Dir: cwd, Jobs: jobs}, + &agenttool.BashOutputTool{Jobs: jobs}, + &agenttool.BashKillTool{Jobs: jobs}, + &agenttool.TodoTool{Store: agenttool.NewTodoStore()}, + &agenttool.WebFetchTool{}, + &agenttool.WebSearchTool{}, + } +} + +// BuiltinToolsExcept returns the default builtin tool set (BuiltinTools) with +// any tool whose name matches one of the except names removed. It backs the +// nesting guard for the generic task tool: a child sub-agent's registry is built +// with "task" excluded so a child can never spawn further sub-agents, capping +// delegation depth at one. With no except names it is equivalent to BuiltinTools. +func BuiltinToolsExcept(cwd string, disabled bool, except ...string) []agentcore.AgentTool { + all := BuiltinTools(cwd, disabled) + if len(except) == 0 || len(all) == 0 { + return all + } + skip := make(map[string]struct{}, len(except)) + for _, n := range except { + skip[n] = struct{}{} + } + out := make([]agentcore.AgentTool, 0, len(all)) + for _, t := range all { + if _, ok := skip[t.Name()]; ok { + continue + } + out = append(out, t) + } + return out +} + +// ReadableExtraRoots returns trusted directories the file tools may reach beyond +// the workspace root. The skills directory is included so the model can load the +// absolute SKILL.md paths pigo advertises in the system prompt, and author or +// update skills there (they otherwise resolve outside the workspace and are +// rejected). An empty skills dir is dropped, so this stays a no-op when the home +// directory cannot be resolved. +func ReadableExtraRoots() []string { + if dir := SkillsDir(); dir != "" { + return []string{dir} + } + return nil +} + +// ToolRegistry builds a registry from the given tools (skipping any that fail to +// register, e.g. a bad schema, which should not happen for built-ins). +func ToolRegistry(tools []agentcore.AgentTool) *agenttool.ToolRegistry { + reg := agenttool.NewToolRegistry() + for _, t := range tools { + _ = reg.Register(t) + } + return reg +} + +// TodoReminders builds the per-turn system-reminder registry for a tool set +// (US-002): it locates the stateful TodoTool and registers a TodoReminderProvider +// over its shared store, so the model is reminded of unfinished tasks each turn. +// It also registers a MemoryReminderProvider over the memory_search tool's store +// (issue #481) when present, so relevant persisted memory is recalled each turn +// (this is the recall channel used after auto-compaction/rebuild). Returns nil +// when neither provider applies (e.g. --no-tools), leaving injection disabled. +func TodoReminders(tools []agentcore.AgentTool) *runtime.ReminderRegistry { + var providers []runtime.ReminderProvider + for _, t := range tools { + switch tool := t.(type) { + case *agenttool.TodoTool: + if tool.Store != nil { + providers = append(providers, &runtime.TodoReminderProvider{Store: tool.Store}) + } + case *agenttool.MemorySearchTool: + if tool.Store != nil { + providers = append(providers, &runtime.MemoryReminderProvider{Store: tool.Store}) + } + } + } + if len(providers) == 0 { + return nil + } + return runtime.NewReminderRegistry(providers...) +} + +// MemoryRootFromTools returns the persistent memory root the run's memory_search +// tool is backed by (its Store.Root()), or "" when persistent memory is not wired +// into this tool set (memory.enabled=false, --no-tools, or the store failed to +// open). It is the canonical source of the memory root for checkpoint persistence +// and context rebuild (/sessions//checkpoint.md): callers resolve the +// root through the opened store rather than re-deriving it from the session store. +func MemoryRootFromTools(tools []agentcore.AgentTool) string { + for _, t := range tools { + if mt, ok := t.(*agenttool.MemorySearchTool); ok && mt.Store != nil { + return mt.Store.Root() + } + } + return "" +} + +// MemoryStoreFromTools returns the persistent memory Store backing the run's +// memory_search tool, or nil when persistent memory is not wired into this tool +// set. It lets status commands (/memory) inspect the live store without +// re-opening the database. +func MemoryStoreFromTools(tools []agentcore.AgentTool) *memory.Store { + for _, t := range tools { + if mt, ok := t.(*agenttool.MemorySearchTool); ok && mt.Store != nil { + return mt.Store + } + } + return nil +} + +// SnapshotRecorderFromTools returns the shared FileSnapshotRecorder backing the +// run's write/edit tools, or nil when file tools are disabled (--no-tools). The +// REPL uses it to commit a per-turn restore point and to serve /rewind. +func SnapshotRecorderFromTools(tools []agentcore.AgentTool) *agenttool.FileSnapshotRecorder { + for _, t := range tools { + switch tool := t.(type) { + case *agenttool.WriteTool: + if tool.Snap != nil { + return tool.Snap + } + case *agenttool.EditTool: + if tool.Snap != nil { + return tool.Snap + } + } + } + return nil +} + +// BashJobStoreFromTools returns the shared BashJobStore backing the run's bash / +// bash_output / kill_bash tools, or nil when the shell tool is disabled. The +// REPL uses it to kill any still-running background jobs on exit so they are not +// orphaned. +func BashJobStoreFromTools(tools []agentcore.AgentTool) *agenttool.BashJobStore { + for _, t := range tools { + if bt, ok := t.(*agenttool.BashTool); ok && bt.Jobs != nil { + return bt.Jobs + } + } + return nil +} + +// MemoryDir returns the persistent memory root directory: $PIGO_HOME/memory, or +// ~/.pigo/memory by default (a single global store so cross-project "global" +// memories are searchable, mirroring the session store's ~/.pigo base). It +// returns "" when the home directory cannot be resolved and no override is set. +func MemoryDir() string { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return "" + } + dir = filepath.Join(home, ".pigo") + } + return filepath.Join(dir, "memory") +} + +// OpenMemoryStore opens the persistent memory store under MemoryDir() (index DB +// at /index.db). It returns (nil, nil) — not an error — when persistent +// memory is disabled (memEnabled=false) or the home dir is unresolvable, so the +// caller degrades to file-based auto-memory without treating the off state as a +// failure. A genuine open failure is returned as an error for the caller to log +// non-fatally. +func OpenMemoryStore(memEnabled bool) (*memory.Store, error) { + if !memEnabled { + return nil, nil + } + root := MemoryDir() + if root == "" { + return nil, nil + } + dbPath := filepath.Join(root, "index.db") + return memory.Open(dbPath, root, "") +} + +// SkillsDir returns the directory skills are loaded from. It defaults to +// ~/.agents/skills, overridable via PIGO_SKILLS_DIR. An empty string is returned +// when the home directory cannot be resolved and no override is set. +func SkillsDir() string { + if dir := os.Getenv("PIGO_SKILLS_DIR"); dir != "" { + return dir + } + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".agents", "skills") +} + +// LoadSkills discovers skills from SkillsDir() once, for both prompt injection +// and /skill-name registration. Under --no-skills it is a no-op. Built-in skills +// are bootstrapped into the skills dir first, then the directory is loaded. +func LoadSkills(noSkills bool) ([]*runtime.Skill, error) { + if noSkills { + return nil, nil + } + var blog io.Writer + if os.Getenv("PIGO_DEBUG") != "" { + blog = os.Stderr + } + builtinskills.Bootstrap(ConfigDir(), SkillsDir(), blog) + dir := SkillsDir() + if dir == "" { + return nil, nil + } + return runtime.LoadSkillsDir(dir) +} + +// PluginsDir returns the directory external plugins are discovered from: +// $PIGO_HOME/plugins, or ~/.pigo/plugins by default. An empty string is returned +// when the home directory cannot be resolved and no override is set (Discover +// then treats it as "no plugins"). +func PluginsDir() string { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return "" + } + dir = filepath.Join(home, ".pigo") + } + return filepath.Join(dir, "plugins") +} + +// ConfigDir returns the directory pigo reads its global config layer from: +// $PIGO_HOME, or ~/.pigo by default. An empty string is returned when the home +// directory cannot be resolved and no override is set (the caller then treats +// the global layer as absent). +func ConfigDir() string { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return "" + } + dir = filepath.Join(home, ".pigo") + } + return dir +} + +// ResolveThinkingLevel resolves the effective reasoning-effort level through the +// layered config chain (US-023): default < global < project < env < CLI flag. +// The global layer is $PIGO_HOME/config.json (or ~/.pigo/config.json); the +// project layer is ./.pigo/config.json in the working directory. A malformed +// layer file or an invalid resolved value is a hard error, surfaced to the +// caller for exit-code mapping. cliLevel is the raw --thinking-level flag ("" = +// unset, so lower layers show through). +func ResolveThinkingLevel(cliLevel string) (agentcore.ThinkingLevel, error) { + def := runtime.DefaultConfigLayer() + layers := []*runtime.ConfigLayer{&def} + + if dir := ConfigDir(); dir != "" { + global, err := runtime.LoadConfigLayer(filepath.Join(dir, "config.json")) + if err != nil { + return "", err + } + layers = append(layers, global) + } + project, err := runtime.LoadConfigLayer(filepath.Join(".pigo", "config.json")) + if err != nil { + return "", err + } + layers = append(layers, project) + + env := runtime.EnvConfigLayer(os.Getenv) + layers = append(layers, &env) + + if v := strings.TrimSpace(cliLevel); v != "" { + cli := runtime.ConfigLayer{ThinkingLevel: &v} + layers = append(layers, &cli) + } + + cfg, err := runtime.ResolveConfig(layers...) + if err != nil { + return "", err + } + return cfg.ThinkingLevel, nil +} + +// ResolveHookSet resolves the effective hook set through the same layered config +// chain as ResolveThinkingLevel (default < global < project < env), with one +// difference required by FR-14: the project layer (./.pigo/config.json under +// cwd) is only merged when the directory is trusted. An untrusted directory +// therefore contributes no hooks, so a checked-out repo cannot run arbitrary +// commands until the user trusts it. A malformed layer file is a hard error, +// surfaced to the caller. The returned set is empty (len 0) when no layer +// defines hooks, which InstallHooks treats as "no hooks" (FR-18). +func ResolveHookSet(cwd string, trusted bool) (hooks.HookSet, error) { + def := runtime.DefaultConfigLayer() + layers := []*runtime.ConfigLayer{&def} + + if dir := ConfigDir(); dir != "" { + global, err := runtime.LoadConfigLayer(filepath.Join(dir, "config.json")) + if err != nil { + return nil, err + } + layers = append(layers, global) + } + if trusted { + project, err := runtime.LoadConfigLayer(filepath.Join(cwd, ".pigo", "config.json")) + if err != nil { + return nil, err + } + layers = append(layers, project) + } + env := runtime.EnvConfigLayer(os.Getenv) + layers = append(layers, &env) + + cfg, err := runtime.ResolveConfig(layers...) + if err != nil { + return nil, err + } + return cfg.Hooks, nil +} + +// Trusted reports whether cwd is a trusted directory per the shared trust store +// ($PIGO_HOME/trust.json). It is the trust gate for the non-interactive drivers +// (headless / TUI / sub-agent) that have no live trust.Manager to consult, so +// ResolveHookSet can honor FR-14 uniformly. A missing or unreadable store is +// treated as untrusted (fail closed): a directory only runs project-layer hooks +// after the user has explicitly trusted it. +func Trusted(cwd string) bool { + m, err := trust.NewManager(trust.DefaultPath()) + if err != nil || m == nil { + return false + } + return m.IsTrusted(cwd) +} + +// NewConfig builds the loop configuration shared by every driver: the provider +// stream, the dynamic API-key resolver, and the tool registry. It is the single +// definition of "how a run is wired", so the REPL (streamRun) and the headless +// driver cannot drift apart. +func NewConfig(model, providerName string, thinking agentcore.ThinkingLevel, prov provider.Provider, creds *provider.CredentialStore, reg *agenttool.ToolRegistry, reminders *runtime.ReminderRegistry) runtime.RunConfig { + return runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: model, + Provider: providerName, + ThinkingLevel: thinking, + Stream: provider.StreamFnFromProvider(prov), + GetAPIKey: creds.GetAPIKey, + }, + Batch: agenttool.BatchConfig{ + ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}, + }, + Reminders: reminders, + } +} diff --git a/pigo/internal/cli/run/task_wiring_test.go b/pigo/internal/cli/run/task_wiring_test.go new file mode 100644 index 0000000..8a2ae79 --- /dev/null +++ b/pigo/internal/cli/run/task_wiring_test.go @@ -0,0 +1,39 @@ +package run + +// Tests for the generic task tool wiring (US-004, #454): the nesting guard. +// BuiltinToolsExcept backs the child sub-agent registry, from which "task" is +// removed so a child cannot spawn further sub-agents. + +import "testing" + +// TestBuiltinToolsExceptExcludesTask verifies the child tool set produced for a +// task sub-agent (builtins minus "task") never contains "task", and that +// excluding a name actually present removes it. +func TestBuiltinToolsExceptExcludesTask(t *testing.T) { + child := BuiltinToolsExcept("/tmp", false, "task") + if len(child) == 0 { + t.Fatal("expected builtin tools in the child set") + } + for _, tl := range child { + if tl.Name() == "task" { + t.Fatal("child tool set must not contain 'task'") + } + } + // Excluding a name that IS present shrinks the set by exactly that tool. + full := BuiltinToolsExcept("/tmp", false) + dropped := BuiltinToolsExcept("/tmp", false, full[0].Name()) + if len(dropped) != len(full)-1 { + t.Errorf("excluding %q: got %d tools, want %d", full[0].Name(), len(dropped), len(full)-1) + } +} + +// TestBuiltinToolsExceptNoExcept verifies that with no except names the result +// matches BuiltinTools, and that a disabled tool set stays empty. +func TestBuiltinToolsExceptNoExcept(t *testing.T) { + if got, want := len(BuiltinToolsExcept("/tmp", false)), len(BuiltinTools("/tmp", false)); got != want { + t.Errorf("no-except size = %d, want %d", got, want) + } + if got := BuiltinToolsExcept("/tmp", true, "task"); got != nil { + t.Errorf("disabled tools should be nil, got %v", got) + } +} diff --git a/pigo/internal/cli/run/thinking_test.go b/pigo/internal/cli/run/thinking_test.go new file mode 100644 index 0000000..c9732b6 --- /dev/null +++ b/pigo/internal/cli/run/thinking_test.go @@ -0,0 +1,120 @@ +package run + +// Tests for resolveThinkingLevel: the CLI end of the layered config chain +// (US-023). It resolves the effective reasoning-effort level through +// default < global ($PIGO_HOME/config.json) < project (./.pigo/config.json) +// < env (PIGO_THINKING_LEVEL) < --thinking-level flag, and rejects an invalid +// value. Each test isolates PIGO_HOME and the working directory so it never +// reads the developer's real config. + +import ( + "os" + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// isolateConfig points PIGO_HOME at a temp dir and chdir's into a temp working +// directory (restored on cleanup), so no real global/project config leaks in. +func isolateConfig(t *testing.T) string { + t.Helper() + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + t.Setenv("PIGO_THINKING_LEVEL", "") + + wd := t.TempDir() + prev, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + if err := os.Chdir(wd); err != nil { + t.Fatalf("chdir: %v", err) + } + t.Cleanup(func() { _ = os.Chdir(prev) }) + return home +} + +// writeConfig writes a config.json layer with the given thinkingLevel at path. +func writeConfig(t *testing.T, path, level string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + body := `{"thinkingLevel":"` + level + `"}` + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } +} + +// TestResolveThinkingLevelDefault verifies the built-in default (medium) applies +// when no layer sets a level. +func TestResolveThinkingLevelDefault(t *testing.T) { + isolateConfig(t) + got, err := ResolveThinkingLevel("") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if got != agentcore.ThinkingMedium { + t.Errorf("level = %q, want medium", got) + } +} + +// TestResolveThinkingLevelFlagWins verifies the --thinking-level flag overrides +// every lower layer (global, project, and env). +func TestResolveThinkingLevelFlagWins(t *testing.T) { + home := isolateConfig(t) + writeConfig(t, filepath.Join(home, "config.json"), "low") + writeConfig(t, filepath.Join(".pigo", "config.json"), "high") + t.Setenv("PIGO_THINKING_LEVEL", "minimal") + + got, err := ResolveThinkingLevel("xhigh") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if got != agentcore.ThinkingXHigh { + t.Errorf("level = %q, want xhigh (flag wins)", got) + } +} + +// TestResolveThinkingLevelEnvOverProject verifies env beats project, and project +// beats global (precedence: global < project < env). +func TestResolveThinkingLevelEnvOverProject(t *testing.T) { + home := isolateConfig(t) + writeConfig(t, filepath.Join(home, "config.json"), "low") + writeConfig(t, filepath.Join(".pigo", "config.json"), "high") + t.Setenv("PIGO_THINKING_LEVEL", "off") + + got, err := ResolveThinkingLevel("") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if got != agentcore.ThinkingOff { + t.Errorf("level = %q, want off (env over project/global)", got) + } +} + +// TestResolveThinkingLevelProjectOverGlobal verifies the project layer overrides +// the global layer when env and flag are unset. +func TestResolveThinkingLevelProjectOverGlobal(t *testing.T) { + home := isolateConfig(t) + writeConfig(t, filepath.Join(home, "config.json"), "low") + writeConfig(t, filepath.Join(".pigo", "config.json"), "high") + + got, err := ResolveThinkingLevel("") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if got != agentcore.ThinkingHigh { + t.Errorf("level = %q, want high (project over global)", got) + } +} + +// TestResolveThinkingLevelInvalid verifies an unknown value is a hard error +// (surfaced for exit-code mapping), not silently coerced. +func TestResolveThinkingLevelInvalid(t *testing.T) { + isolateConfig(t) + if _, err := ResolveThinkingLevel("turbo"); err == nil { + t.Error("expected error for invalid thinking level, got nil") + } +} diff --git a/pigo/internal/cli/run/toolpolicy.go b/pigo/internal/cli/run/toolpolicy.go new file mode 100644 index 0000000..15335f5 --- /dev/null +++ b/pigo/internal/cli/run/toolpolicy.go @@ -0,0 +1,214 @@ +package run + +import ( + "fmt" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// ToolPolicy is the user-declared tool boundary for a run: the --allowed-tools +// whitelist and the --disallowed-tools blacklist, already normalized by +// SplitToolNames. It is passed as one value rather than two adjacent []string +// parameters so the two lists cannot be swapped at a call site — silently +// inverting a security boundary is exactly the bug that must be impossible. +// +// The zero value means "no restriction" and every operation on it is a no-op. +type ToolPolicy struct { + Allow []string + Deny []string +} + +// NewToolPolicy normalizes raw flag values into a policy. +func NewToolPolicy(allowed, disallowed []string) ToolPolicy { + return ToolPolicy{Allow: SplitToolNames(allowed), Deny: SplitToolNames(disallowed)} +} + +// IsZero reports whether the policy constrains nothing. +func (p ToolPolicy) IsZero() bool { return len(p.Allow) == 0 && len(p.Deny) == 0 } + +// SplitToolNames normalizes raw --allowed-tools / --disallowed-tools values into +// a flat list of tool names. Each value may itself be a comma-separated list, so +// `--allowed-tools "read,grep"` and `--allowed-tools read --allowed-tools grep` +// are equivalent. Entries are lowercased and trimmed, and empty entries dropped, +// so `"read, ,grep"` yields [read grep]. Matching is case-insensitive on purpose: +// users coming from Claude Code write `Read`/`Bash`, which must hit pigo's +// `read`/`bash`. +func SplitToolNames(values []string) []string { + if len(values) == 0 { + return nil + } + out := make([]string, 0, len(values)) + for _, v := range values { + for _, part := range strings.Split(v, ",") { + if n := normalizeToolName(part); n != "" { + out = append(out, n) + } + } + } + if len(out) == 0 { + return nil + } + return out +} + +// normalizeToolName is the single definition of how a tool name is compared: +// surrounding whitespace is insignificant and case is ignored. +// +// Note this is deliberately more lenient than runtime.filterToolsByName, which +// backs a skill frontmatter's allowed-tools and matches case-sensitively. The +// two are not unified: this policy is user-facing CLI input where Claude-Code +// habits (Read/Bash) must hit read/bash, whereas skill frontmatter is authored +// against pigo's canonical lowercase names. Keep them separate on purpose. +func normalizeToolName(s string) string { + return strings.ToLower(strings.TrimSpace(s)) +} + +// ToolPolicyError reports --allowed-tools / --disallowed-tools entries that name +// no existing tool. It is a usage error — the caller maps it to exit code 2 — +// because silently ignoring a typo is the worst outcome available: the user +// believes a boundary is in force when it is not. +type ToolPolicyError struct { + // UnknownAllowed and UnknownDisallowed are the unrecognized names from each + // flag. Both are reported in one error so a user with two typos fixes both in + // one round rather than one per run. + UnknownAllowed []string + UnknownDisallowed []string + // Available is the sorted set of names that would have been accepted. + Available []string +} + +func (e *ToolPolicyError) Error() string { + var parts []string + if len(e.UnknownAllowed) > 0 { + parts = append(parts, fmt.Sprintf("--allowed-tools: unknown tool %s", quoteNames(e.UnknownAllowed))) + } + if len(e.UnknownDisallowed) > 0 { + parts = append(parts, fmt.Sprintf("--disallowed-tools: unknown tool %s", quoteNames(e.UnknownDisallowed))) + } + return fmt.Sprintf("%s (available: %s)", strings.Join(parts, "; "), strings.Join(e.Available, ", ")) +} + +// quoteNames renders names as a comma-separated quoted list. +func quoteNames(names []string) string { + out := make([]string, len(names)) + for i, n := range names { + out[i] = fmt.Sprintf("%q", n) + } + return strings.Join(out, ", ") +} + +// ValidateToolPolicy checks every allow/deny entry against the assembled tool +// set. It must run AFTER the full set exists (builtins + memory + task + +// plugins), because plugin and memory tool names are only known at runtime; +// validating right after flag parsing would reject legitimate plugin names. +// +// An empty tool set (--no-tools) skips validation entirely: there is nothing to +// constrain, and reporting every name as unknown would be noise. +func ValidateToolPolicy(tools []agentcore.AgentTool, policy ToolPolicy) error { + if len(tools) == 0 || policy.IsZero() { + return nil + } + known := toolNameSet(tools) + unknownAllow := unknownNames(known, policy.Allow) + unknownDeny := unknownNames(known, policy.Deny) + if len(unknownAllow) == 0 && len(unknownDeny) == 0 { + return nil + } + available := make([]string, 0, len(known)) + for n := range known { + available = append(available, n) + } + sort.Strings(available) + return &ToolPolicyError{ + UnknownAllowed: unknownAllow, + UnknownDisallowed: unknownDeny, + Available: available, + } +} + +// unknownNames returns the entries of names absent from known, preserving input +// order and dropping duplicates so a name repeated twice is reported once. +func unknownNames(known map[string]struct{}, names []string) []string { + var out []string + seen := make(map[string]struct{}, len(names)) + for _, n := range names { + if _, ok := known[n]; ok { + continue + } + if _, dup := seen[n]; dup { + continue + } + seen[n] = struct{}{} + out = append(out, n) + } + return out +} + +// toolNameSet indexes a tool set by normalized name. +func toolNameSet(tools []agentcore.AgentTool) map[string]struct{} { + set := make(map[string]struct{}, len(tools)) + for _, t := range tools { + set[normalizeToolName(t.Name())] = struct{}{} + } + return set +} + +// ApplyToolPolicy narrows a tool set to the allow list and then removes the deny +// list. Both empty means no restriction and the input is returned unchanged, so +// the default path is a true no-op. +// +// Deny runs after allow, which makes deny win when a name appears on both sides. +// That ordering is deliberate: the fail-closed reading of a contradictory policy +// is "do not run it". +// +// This filters the set handed to the model, so it sits at the tool-registration +// layer — strictly before the BeforeToolCall confirmation gate. A removed tool is +// never advertised and never dispatchable, which is why --approve (which only +// waives per-call confirmation) cannot widen the boundary. +func ApplyToolPolicy(tools []agentcore.AgentTool, policy ToolPolicy) []agentcore.AgentTool { + if len(tools) == 0 || policy.IsZero() { + return tools + } + allowSet := nameSet(policy.Allow) + denySet := nameSet(policy.Deny) + out := make([]agentcore.AgentTool, 0, len(tools)) + for _, t := range tools { + name := normalizeToolName(t.Name()) + if len(allowSet) > 0 { + if _, ok := allowSet[name]; !ok { + continue + } + } + if _, denied := denySet[name]; denied { + continue + } + out = append(out, t) + } + return out +} + +// nameSet indexes already-normalized names for lookup. +func nameSet(names []string) map[string]struct{} { + if len(names) == 0 { + return nil + } + set := make(map[string]struct{}, len(names)) + for _, n := range names { + set[n] = struct{}{} + } + return set +} + +// ChildToolSet builds the tool set for a task sub-agent: the builtins with +// "task" removed (the nesting guard capping delegation depth at one), narrowed by +// the parent's policy. +// +// Inheriting the policy is load-bearing, not a nicety. A child that ignored it +// would be a one-line escape from the boundary: under --disallowed-tools bash the +// model could dispatch a sub-agent and run bash there instead. Any future spawn +// path must route through here for the same reason. +func ChildToolSet(cwd string, policy ToolPolicy) []agentcore.AgentTool { + return ApplyToolPolicy(BuiltinToolsExcept(cwd, false, "task"), policy) +} diff --git a/pigo/internal/cli/run/toolpolicy_setup_test.go b/pigo/internal/cli/run/toolpolicy_setup_test.go new file mode 100644 index 0000000..24dd6b9 --- /dev/null +++ b/pigo/internal/cli/run/toolpolicy_setup_test.go @@ -0,0 +1,212 @@ +package run + +import ( + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" +) + +// writePolicySkill drops a minimal valid skill into dir so LoadSkills has +// something to advertise. +func writePolicySkill(t *testing.T, dir, name, description string) { + t.Helper() + body := "---\nname: " + name + "\ndescription: " + description + "\n---\nDo the thing." + if err := os.WriteFile(filepath.Join(dir, name+".md"), []byte(body), 0o644); err != nil { + t.Fatalf("write skill %s: %v", name, err) + } +} + +// setupToolNames runs SetupEnv with a policy and returns the resulting tool +// names. The provider is never contacted, so a stub model id is fine; --no-skills +// keeps the run independent of the machine's skills directory. +func setupToolNames(t *testing.T, policy ToolPolicy) []string { + t.Helper() + t.Setenv("OPENROUTER_API_KEY", "test-key") + t.Setenv("PIGO_HOME", t.TempDir()) // isolate plugin/skill discovery + env, err := SetupEnv("openrouter/free", "", "", "", "", false /*noTools*/, true /*noSkills*/, "", nil, false /*memEnabled*/, policy) + if err != nil { + t.Fatalf("SetupEnv: %v", err) + } + return names(env.Tools) +} + +// contains reports whether name is in the set. +func contains(set []string, name string) bool { + for _, n := range set { + if n == name { + return true + } + } + return false +} + +// TestSetupEnvAppliesAllowList confirms a whitelist narrows the advertised set. +// The `task` tool is expected to survive only when explicitly allowed. +func TestSetupEnvAppliesAllowList(t *testing.T) { + got := setupToolNames(t, NewToolPolicy([]string{"read,grep"}, nil)) + want := []string{"read", "grep"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Errorf("tool set = %q, want exactly %q", got, want) + } +} + +// TestSetupEnvAppliesDenyList confirms a blacklist removes the named tools while +// leaving everything else — including the side-effect tools not named — in place. +func TestSetupEnvAppliesDenyList(t *testing.T) { + got := setupToolNames(t, NewToolPolicy(nil, []string{"bash", "bash_output", "kill_bash"})) + for _, denied := range []string{"bash", "bash_output", "kill_bash"} { + if contains(got, denied) { + t.Errorf("%q survived the deny list: %q", denied, got) + } + } + for _, kept := range []string{"read", "write", "edit", "grep"} { + if !contains(got, kept) { + t.Errorf("%q was removed but was not denied: %q", kept, got) + } + } +} + +// TestSetupEnvDenyWinsOverAllow is the fail-closed guarantee: a tool named on +// both sides is removed. +func TestSetupEnvDenyWinsOverAllow(t *testing.T) { + got := setupToolNames(t, NewToolPolicy([]string{"read", "bash"}, []string{"bash"})) + if contains(got, "bash") { + t.Errorf("bash was on both lists and must be removed, got %q", got) + } + if !contains(got, "read") { + t.Errorf("read was allowed and not denied, so it must survive, got %q", got) + } +} + +// TestSetupEnvUnconstrainedIsUnchanged is the zero-regression check: no policy +// means the full built-in set, including the side-effect tools. +func TestSetupEnvUnconstrainedIsUnchanged(t *testing.T) { + got := setupToolNames(t, ToolPolicy{}) + for _, want := range []string{"read", "write", "edit", "grep", "find", "bash", "todo", "webfetch", "websearch", "task"} { + if !contains(got, want) { + t.Errorf("unconstrained run is missing %q: %q", want, got) + } + } +} + +// TestSetupEnvRejectsUnknownToolName confirms a typo aborts setup with a +// ToolPolicyError, which is what maps to exit code 2 rather than a run that +// silently ignores the boundary. +func TestSetupEnvRejectsUnknownToolName(t *testing.T) { + t.Setenv("OPENROUTER_API_KEY", "test-key") + t.Setenv("PIGO_HOME", t.TempDir()) + _, err := SetupEnv("openrouter/free", "", "", "", "", false, true, "", nil, false, NewToolPolicy([]string{"raed"}, nil)) + if err == nil { + t.Fatal("SetupEnv = nil error, want a failure for the misspelled tool name") + } + var policyErr *ToolPolicyError + if !errors.As(err, &policyErr) { + t.Fatalf("error type = %T, want *ToolPolicyError", err) + } +} + +// TestChildToolSetInheritsPolicy closes the sub-agent escape hatch: a child +// dispatched by the task tool must not regain a tool the parent's policy removed, +// or `--disallowed-tools bash` would be bypassable by delegating. +func TestChildToolSetInheritsPolicy(t *testing.T) { + child := names(ChildToolSet("/tmp", NewToolPolicy(nil, []string{"bash"}))) + if contains(child, "bash") { + t.Errorf("child regained the denied bash tool: %q", child) + } + if contains(child, "task") { + t.Errorf("child must not contain task (nesting guard): %q", child) + } + if !contains(child, "read") { + t.Errorf("child lost an un-denied tool: %q", child) + } + + allowOnly := names(ChildToolSet("/tmp", NewToolPolicy([]string{"read"}, nil))) + if strings.Join(allowOnly, ",") != "read" { + t.Errorf("child under an allow list = %q, want exactly [read]", allowOnly) + } + + // With no policy the child is the plain nesting-guarded builtin set. + unconstrained := names(ChildToolSet("/tmp", ToolPolicy{})) + if !contains(unconstrained, "bash") || contains(unconstrained, "task") { + t.Errorf("unconstrained child set = %q, want builtins minus task", unconstrained) + } +} + +// TestSetupEnvSkillsGatedOnFilteredReadTool covers the ordering dependency: the +// block is advertised only when `read` survives the policy, +// because the model needs read to load a skill body. Filtering must therefore +// happen before the system prompt is built. +func TestSetupEnvSkillsGatedOnFilteredReadTool(t *testing.T) { + t.Setenv("OPENROUTER_API_KEY", "test-key") + t.Setenv("PIGO_HOME", t.TempDir()) + skillsDir := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", skillsDir) + writePolicySkill(t, skillsDir, "weather", "get the weather") + + withRead, err := SetupEnv("openrouter/free", "", "", "", "", false, false, "", nil, false, ToolPolicy{}) + if err != nil { + t.Fatalf("SetupEnv (unconstrained): %v", err) + } + if !strings.Contains(withRead.SysPrompt, "") { + t.Fatal("unconstrained run must advertise skills; the fixture or gate is wrong") + } + + withoutRead, err := SetupEnv("openrouter/free", "", "", "", "", false, false, "", nil, false, NewToolPolicy(nil, []string{"read"})) + if err != nil { + t.Fatalf("SetupEnv (read denied): %v", err) + } + if strings.Contains(withoutRead.SysPrompt, "available_skills") { + t.Error("denying read must suppress : the model could not load a skill body") + } +} + +// captureStderr runs fn with os.Stderr redirected to a pipe and returns whatever +// was written. It is not safe under t.Parallel — these tests must stay serial. +func captureStderr(t *testing.T, fn func()) string { + t.Helper() + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("os.Pipe: %v", err) + } + orig := os.Stderr + os.Stderr = w + defer func() { os.Stderr = orig }() + fn() + w.Close() + out, err := io.ReadAll(r) + if err != nil { + t.Fatalf("read captured stderr: %v", err) + } + return string(out) +} + +// TestSetupEnvNoToolsWithPolicyWarns is the counterpart to the typo guarantee: +// under --no-tools the set is empty, so ValidateToolPolicy cannot flag a +// misspelled name. Rather than let the boundary silently vanish, SetupEnv must +// still succeed but print a warning that the policy is inert — otherwise a user +// combining --no-tools with a (possibly misspelled) --allowed-tools would +// believe a boundary is in force when none is. +func TestSetupEnvNoToolsWithPolicyWarns(t *testing.T) { + t.Setenv("OPENROUTER_API_KEY", "test-key") + t.Setenv("PIGO_HOME", t.TempDir()) + + var env Env + var err error + stderr := captureStderr(t, func() { + // A deliberately misspelled name: with tools present this would abort with + // exit code 2, but --no-tools skips validation, so it must not error. + env, err = SetupEnv("openrouter/free", "", "", "", "", true /*noTools*/, true /*noSkills*/, "", nil, false, NewToolPolicy([]string{"raed"}, nil)) + }) + if err != nil { + t.Fatalf("SetupEnv(--no-tools + policy) = %v, want nil (validation is skipped, not failed)", err) + } + if len(env.Tools) != 0 { + t.Errorf("--no-tools must leave no tools, got %q", names(env.Tools)) + } + if !strings.Contains(stderr, "--no-tools disables all tools") { + t.Errorf("expected an inert-policy warning on stderr, got %q", stderr) + } +} diff --git a/pigo/internal/cli/run/toolpolicy_test.go b/pigo/internal/cli/run/toolpolicy_test.go new file mode 100644 index 0000000..35aa608 --- /dev/null +++ b/pigo/internal/cli/run/toolpolicy_test.go @@ -0,0 +1,176 @@ +package run + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// policyTool is a minimal AgentTool whose only meaningful property is its name — +// the tool policy matches on nothing else. +type policyTool struct{ name string } + +func (t policyTool) Name() string { return t.name } +func (t policyTool) Description() string { return "stub" } +func (t policyTool) Schema() json.RawMessage { + return json.RawMessage(`{"type":"object"}`) +} +func (t policyTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} +func (t policyTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{}, nil +} + +// toolSet builds a tool set from names. +func toolSet(names ...string) []agentcore.AgentTool { + out := make([]agentcore.AgentTool, 0, len(names)) + for _, n := range names { + out = append(out, policyTool{name: n}) + } + return out +} + +// names extracts the tool names from a set, for comparison. +func names(tools []agentcore.AgentTool) []string { + out := make([]string, 0, len(tools)) + for _, t := range tools { + out = append(out, t.Name()) + } + return out +} + +// TestSplitToolNames covers the accepted input forms: a single value, repeated +// flags, comma-separated values, mixed forms, and whitespace/empty entries. +func TestSplitToolNames(t *testing.T) { + tests := []struct { + name string + in []string + want []string + }{ + {"nil", nil, nil}, + {"single", []string{"read"}, []string{"read"}}, + {"repeated flag", []string{"read", "grep"}, []string{"read", "grep"}}, + {"comma", []string{"read,grep"}, []string{"read", "grep"}}, + {"mixed", []string{"read,grep", "bash"}, []string{"read", "grep", "bash"}}, + {"whitespace and empties", []string{"read, ,grep", " bash "}, []string{"read", "grep", "bash"}}, + {"case folded", []string{"Read", "BASH"}, []string{"read", "bash"}}, + {"only empties", []string{"", " ", ",,"}, nil}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := SplitToolNames(tc.in) + if len(got) != len(tc.want) { + t.Fatalf("SplitToolNames(%q) = %q, want %q", tc.in, got, tc.want) + } + for i := range got { + if got[i] != tc.want[i] { + t.Fatalf("SplitToolNames(%q) = %q, want %q", tc.in, got, tc.want) + } + } + }) + } +} + +// TestApplyToolPolicy covers allow-only, deny-only, the overlap (deny wins), the +// unconstrained no-op, and filtering down to an empty set. +func TestApplyToolPolicy(t *testing.T) { + all := toolSet("read", "write", "bash", "grep") + tests := []struct { + name string + policy ToolPolicy + want []string + }{ + {"unconstrained is a no-op", ToolPolicy{}, []string{"read", "write", "bash", "grep"}}, + {"allow only", NewToolPolicy([]string{"read,grep"}, nil), []string{"read", "grep"}}, + {"deny only", NewToolPolicy(nil, []string{"bash"}), []string{"read", "write", "grep"}}, + {"deny wins over allow", NewToolPolicy([]string{"read,bash"}, []string{"bash"}), []string{"read"}}, + {"case-insensitive allow", NewToolPolicy([]string{"Read", "GREP"}, nil), []string{"read", "grep"}}, + {"case-insensitive deny", NewToolPolicy(nil, []string{"Bash"}), []string{"read", "write", "grep"}}, + {"filters to empty", NewToolPolicy([]string{"read"}, []string{"read"}), nil}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := names(ApplyToolPolicy(all, tc.policy)) + if strings.Join(got, ",") != strings.Join(tc.want, ",") { + t.Errorf("ApplyToolPolicy = %q, want %q", got, tc.want) + } + }) + } +} + +// TestApplyToolPolicyEmptyToolSet confirms an empty input (--no-tools) is left +// alone rather than being treated as "everything denied". +func TestApplyToolPolicyEmptyToolSet(t *testing.T) { + if got := ApplyToolPolicy(nil, NewToolPolicy([]string{"read"}, nil)); got != nil { + t.Errorf("ApplyToolPolicy(nil, ...) = %v, want nil", got) + } +} + +// TestValidateToolPolicyAccepts confirms known names — including case variants — +// pass validation. +func TestValidateToolPolicyAccepts(t *testing.T) { + all := toolSet("read", "bash") + for _, policy := range []ToolPolicy{ + {}, + NewToolPolicy([]string{"read"}, nil), + NewToolPolicy([]string{"Read"}, []string{"BASH"}), + NewToolPolicy(nil, []string{"bash"}), + } { + if err := ValidateToolPolicy(all, policy); err != nil { + t.Errorf("ValidateToolPolicy(%+v) = %v, want nil", policy, err) + } + } +} + +// TestValidateToolPolicyReportsAllUnknown is the anti-typo guarantee: every bad +// name from both flags is reported in one error, alongside the available names, +// so a user with two typos fixes both in one round. +func TestValidateToolPolicyReportsAllUnknown(t *testing.T) { + all := toolSet("read", "bash", "grep") + err := ValidateToolPolicy(all, NewToolPolicy([]string{"raed,gerp"}, []string{"bahs"})) + if err == nil { + t.Fatal("ValidateToolPolicy = nil, want an error for the misspelled names") + } + var policyErr *ToolPolicyError + if !errors.As(err, &policyErr) { + t.Fatalf("error type = %T, want *ToolPolicyError (exit-code mapping depends on it)", err) + } + if len(policyErr.UnknownAllowed) != 2 { + t.Errorf("UnknownAllowed = %q, want both misspellings", policyErr.UnknownAllowed) + } + if len(policyErr.UnknownDisallowed) != 1 { + t.Errorf("UnknownDisallowed = %q, want the one misspelling", policyErr.UnknownDisallowed) + } + msg := err.Error() + for _, want := range []string{`"raed"`, `"gerp"`, `"bahs"`, "available:", "read", "bash", "grep"} { + if !strings.Contains(msg, want) { + t.Errorf("error message %q is missing %q", msg, want) + } + } +} + +// TestValidateToolPolicyDeduplicatesUnknown confirms a name repeated across +// values is reported once. +func TestValidateToolPolicyDeduplicatesUnknown(t *testing.T) { + err := ValidateToolPolicy(toolSet("read"), NewToolPolicy([]string{"raed", "raed"}, nil)) + var policyErr *ToolPolicyError + if !errors.As(err, &policyErr) { + t.Fatalf("error = %v, want *ToolPolicyError", err) + } + if len(policyErr.UnknownAllowed) != 1 { + t.Errorf("UnknownAllowed = %q, want one entry", policyErr.UnknownAllowed) + } +} + +// TestValidateToolPolicySkipsEmptyToolSet confirms --no-tools does not turn every +// policy name into an error. +func TestValidateToolPolicySkipsEmptyToolSet(t *testing.T) { + if err := ValidateToolPolicy(nil, NewToolPolicy([]string{"anything"}, nil)); err != nil { + t.Errorf("ValidateToolPolicy(nil tools) = %v, want nil", err) + } +} diff --git a/pigo/internal/cli/status/fakehost_test.go b/pigo/internal/cli/status/fakehost_test.go new file mode 100644 index 0000000..efc1bd1 --- /dev/null +++ b/pigo/internal/cli/status/fakehost_test.go @@ -0,0 +1,49 @@ +package status + +// fakeHost satisfies cli.Host by embedding the interface (so every method is +// present) while overriding only the accessors RunStatus reads: Live, Header, +// AgentCtx, Cwd, Trust, Slash, Creds, Telemetry. The embedded nil interface +// would panic if any other method were called, which these tests never do. +// This lets the status tests run without the package-main REPL harness. + +import ( + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +type fakeHost struct { + cli.Host + live *cli.LiveConfig + header session.SessionHeader + agentCtx *agentcore.AgentContext + cwd string + trust *trust.Manager + slash *runtime.SlashRegistry + creds *provider.CredentialStore + telemetry *cli.TelemetryHolder +} + +func (f *fakeHost) Live() *cli.LiveConfig { return f.live } +func (f *fakeHost) Header() session.SessionHeader { return f.header } +func (f *fakeHost) AgentCtx() *agentcore.AgentContext { return f.agentCtx } +func (f *fakeHost) Cwd() string { return f.cwd } +func (f *fakeHost) Trust() *trust.Manager { return f.trust } +func (f *fakeHost) Slash() *runtime.SlashRegistry { return f.slash } +func (f *fakeHost) Creds() *provider.CredentialStore { return f.creds } +func (f *fakeHost) Telemetry() *cli.TelemetryHolder { return f.telemetry } + +// newFakeHost builds a fakeHost with empty-but-non-nil live config, agent +// context, slash registry and credential store, mirroring a fresh session. +// Tests customize the returned host's fields before calling RunStatus. +func newFakeHost() *fakeHost { + return &fakeHost{ + live: &cli.LiveConfig{}, + agentCtx: &agentcore.AgentContext{}, + slash: runtime.NewSlashRegistry(), + creds: provider.NewCredentialStore(nil), + } +} diff --git a/pigo/internal/cli/status/status.go b/pigo/internal/cli/status/status.go new file mode 100644 index 0000000..d0a7274 --- /dev/null +++ b/pigo/internal/cli/status/status.go @@ -0,0 +1,292 @@ +// This file implements the /status slash command (US-002, #292) that prints a +// colored multi-section status report with runtime config, context usage, and more. +// +// It reaches the session's collaborators and mutable state through the cli.Host +// contract (like /goal and /btw) rather than importing the concrete replDeps +// aggregate, keeping the dependency single-direction (repl→status). +package status + +import ( + "context" + "fmt" + "io" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/trust" +) + +// RunStatus prints a colored multi-section status report to out using data +// read from host through the cli.Host accessors. It shows runtime config, +// context usage, project/environment, credentials, and telemetry. +func RunStatus(out io.Writer, host cli.Host) { + color := ui.Enabled() + + fmt.Fprintln(out) + printRuntimeConfig(out, color, host) + fmt.Fprintln(out) + printContextStatus(out, color, host) + fmt.Fprintln(out) + printEnvStatus(out, color, host) + fmt.Fprintln(out) + printCredentialsStatus(out, color, host) + fmt.Fprintln(out) + printTelemetryStatus(out, color, host) +} + +// printRuntimeConfig prints the runtime model configuration section. +func printRuntimeConfig(out io.Writer, color bool, host cli.Host) { + live := host.Live() + header := host.Header() + model := live.Model + providerName := live.ProviderName + baseURL := live.BaseURL + protocol := live.Protocol + thinkingLevel := string(live.ThinkingLevel) + contextWindow := live.ContextWindow + + if model == "" { + model = header.Model + } + if providerName == "" { + providerName = header.Provider + } + if baseURL == "" { + baseURL = "(default)" + } + if protocol == "" { + protocol = "(default)" + } + if thinkingLevel == "" { + thinkingLevel = "(default)" + } + + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "runtime config:")) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "model:"), model) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "provider:"), providerName) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "base URL:"), baseURL) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "protocol:"), protocol) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "thinking:"), thinkingLevel) + if contextWindow > 0 { + fmt.Fprintf(out, " %s %d tokens\n", ui.Colorize(color, ui.Dim, "context window:"), contextWindow) + } else { + fmt.Fprintf(out, " %s unknown\n", ui.Colorize(color, ui.Dim, "context window:")) + } +} + +// printContextStatus prints the current context usage and compaction section. +func printContextStatus(out io.Writer, color bool, host cli.Host) { + msgs := host.AgentCtx().Messages + tokens := compaction.EstimateContextTokens(msgs).Tokens + contextWindow := host.Live().ContextWindow + + compactions := 0 + for _, m := range msgs { + if _, ok := m.(agentcore.CompactionMessage); ok { + compactions++ + } + } + + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "context:")) + fmt.Fprintf(out, " %s %d / %d tokens\n", ui.Colorize(color, ui.Dim, "current:"), tokens, contextWindow) + + // Calculate utilization percentage if possible + if contextWindow > 0 { + utilization := int(float64(tokens) / float64(contextWindow) * 100) + utilColor := "" + if utilization >= 90 { + utilColor = ui.Red + } else if utilization >= 70 { + utilColor = ui.Yellow + } else { + utilColor = ui.Green + } + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "utilization:"), ui.Colorize(color, utilColor, fmt.Sprintf("%d%%", utilization))) + } else { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "utilization:"), ui.Colorize(color, ui.Yellow, "unknown")) + } + + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "compactions:"), compactions) + + // Calculate remaining tokens before auto-compaction + if contextWindow > 0 { + reserve := compaction.DefaultCompactionSettings.ReserveTokens + threshold := contextWindow - reserve + remaining := threshold - tokens + if remaining < 0 { + fmt.Fprintf(out, " %s %s (threshold: %d, reserve: %d)\n", + ui.Colorize(color, ui.Dim, "before compact:"), + ui.Colorize(color, ui.Red, fmt.Sprintf("%d over threshold", -remaining)), + threshold, + reserve, + ) + } else { + fmt.Fprintf(out, " %s %s (threshold: %d, reserve: %d)\n", + ui.Colorize(color, ui.Dim, "before compact:"), + ui.Colorize(color, ui.Green, fmt.Sprintf("%d tokens remaining", remaining)), + threshold, + reserve, + ) + } + } else { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "before compact:"), + ui.Colorize(color, ui.Yellow, "auto-compaction disabled (unknown window)"), + ) + } +} + +// printEnvStatus prints the project & environment section: cwd, trust status, +// and counts of loaded skills and plugins (with names). User command templates +// are listed separately when present. +func printEnvStatus(out io.Writer, color bool, host cli.Host) { + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "project & environment:")) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "cwd:"), host.Cwd()) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "trust:"), trustStatus(host.Trust(), host.Cwd())) + + var skills, plugins, userCmds []string + if slash := host.Slash(); slash != nil { + for _, c := range slash.List() { + switch c.Source { + case runtime.SourceSkill: + skills = append(skills, c.Name) + case runtime.SourcePlugin: + plugins = append(plugins, c.Name) + case runtime.SourceUser: + userCmds = append(userCmds, c.Name) + } + } + } + fmt.Fprintf(out, " %s %d%s\n", ui.Colorize(color, ui.Dim, "skills:"), len(skills), namesSuffix(skills)) + fmt.Fprintf(out, " %s %d%s\n", ui.Colorize(color, ui.Dim, "plugins:"), len(plugins), namesSuffix(plugins)) + if len(userCmds) > 0 { + fmt.Fprintf(out, " %s %d%s\n", ui.Colorize(color, ui.Dim, "user commands:"), len(userCmds), namesSuffix(userCmds)) + } +} + +// trustStatus classifies the cwd's trust state for display: disabled when trust +// is off, trusted when IsTrusted is true (session grant or saved Trusted), +// untrusted when a saved Untrusted decision applies, else prompt (undecided). +func trustStatus(mgr *trust.Manager, cwd string) string { + if mgr == nil { + return "disabled" + } + if mgr.IsTrusted(cwd) { + return "trusted" + } + if res := mgr.NearestTrustDecision(cwd); res.Found && res.Decision == trust.Untrusted { + return "untrusted" + } + return "prompt" +} + +// namesSuffix renders " (n1, n2, ...)" for a non-empty name list, capping at 8 +// names with "+k more". It returns "" for an empty list. +func namesSuffix(names []string) string { + if len(names) == 0 { + return "" + } + sort.Strings(names) + const max = 8 + if len(names) <= max { + return " (" + strings.Join(names, ", ") + ")" + } + return fmt.Sprintf(" (%s, +%d more)", strings.Join(names[:max], ", "), len(names)-max) +} + +// printCredentialsStatus prints the credentials & connectivity section: API key +// presence (masked, never plaintext) and the provider endpoint URL. +func printCredentialsStatus(out io.Writer, color bool, host cli.Host) { + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "credentials & connectivity:")) + live := host.Live() + creds := host.Creds() + provider := live.ProviderName + if creds != nil && creds.HasCredential(context.Background(), provider) { + key := creds.GetAPIKey(context.Background(), provider) + fmt.Fprintf(out, " %s %s %s\n", + ui.Colorize(color, ui.Dim, "api key:"), + ui.Colorize(color, ui.Green, "set"), + ui.Colorize(color, ui.Dim, maskKey(key))) + } else { + fmt.Fprintf(out, " %s %s\n", + ui.Colorize(color, ui.Dim, "api key:"), + ui.Colorize(color, ui.Yellow, "not set")) + } + endpoint := live.BaseURL + if endpoint == "" { + endpoint = "(default)" + } + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "endpoint:"), endpoint) +} + +// maskKey returns a masked hint of an API key showing only the last 4 chars +// (e.g. "••••abcd"). A key of 4 chars or fewer is masked entirely. It never +// returns the full key. +func maskKey(key string) string { + const tail = 4 + r := []rune(key) + if len(r) <= tail { + return strings.Repeat("•", len(r)) + } + return strings.Repeat("•", tail) + string(r[len(r)-tail:]) +} + +// printTelemetryStatus prints the telemetry section with two sub-blocks - +// cumulative (since session start) and last run - each showing turn count, +// truncation/compaction counts, context utilization, and a per-tool table. +// Both blocks show "no telemetry yet" before any run has completed. +func printTelemetryStatus(out io.Writer, color bool, host cli.Host) { + fmt.Fprintf(out, "%s\n", ui.Colorize(color, ui.Bold, "telemetry:")) + holder := host.Telemetry() + if holder == nil || !holder.HasTelemetry() { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "since session start:"), ui.Colorize(color, ui.Dim, "no telemetry yet")) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "last run:"), ui.Colorize(color, ui.Dim, "no telemetry yet")) + return + } + printTelemetryBlock(out, color, "since session start:", + holder.CumulativeTurns(), + holder.CumulativeTruncationCount(), + holder.CumulativeCompactionCount(), + holder.CumulativeContextUtilization(), + holder.CumulativeToolDurations(), + ) + if last := holder.Last(); last != nil { + printTelemetryBlock(out, color, "last run:", + last.Turns, + last.TruncationCount, + last.CompactionCount, + last.ContextUtilization, + last.ToolDurationsMs, + ) + } else { + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "last run:"), ui.Colorize(color, ui.Dim, "no telemetry yet")) + } +} + +// printTelemetryBlock renders one telemetry sub-block (cumulative or last run). +func printTelemetryBlock(out io.Writer, color bool, label string, turns, trunc, compact int, util float64, tools map[string]agentcore.ToolTiming) { + fmt.Fprintf(out, " %s\n", ui.Colorize(color, ui.Cyan, label)) + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "turns:"), turns) + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "truncations:"), trunc) + fmt.Fprintf(out, " %s %d\n", ui.Colorize(color, ui.Dim, "compactions:"), compact) + fmt.Fprintf(out, " %s %s\n", ui.Colorize(color, ui.Dim, "utilization:"), fmt.Sprintf("%.0f%%", util*100)) + if len(tools) == 0 { + fmt.Fprintf(out, " %s (none)\n", ui.Colorize(color, ui.Dim, "tools:")) + return + } + names := make([]string, 0, len(tools)) + for n := range tools { + names = append(names, n) + } + sort.Strings(names) + fmt.Fprintf(out, " %s\n", ui.Colorize(color, ui.Dim, "tools:")) + for _, n := range names { + t := tools[n] + fmt.Fprintf(out, " %-12s %3d calls %dms\n", n, t.Count, t.TotalMs) + } +} diff --git a/pigo/internal/cli/status/status_e2e_test.go b/pigo/internal/cli/status/status_e2e_test.go new file mode 100644 index 0000000..bb61ef5 --- /dev/null +++ b/pigo/internal/cli/status/status_e2e_test.go @@ -0,0 +1,139 @@ +// This file holds the end-to-end and edge-case tests for /status (US-005, #295): +// fresh-session behavior, model-switch reflection, telemetry reset on session +// switch, and render timing, exercised through direct RunStatus calls with a +// fake host. The REPL-intercept and headless-flag tests live in package main. +package status + +import ( + "bytes" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" +) + +// TestStatusFreshSessionAllSections verifies /status renders every section on a +// brand-new session before any model turn, with "no telemetry yet", and does not +// panic. +func TestStatusFreshSessionAllSections(t *testing.T) { + host := newFakeHost() + host.cwd = "/tmp/e2e-fresh" + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + for _, want := range []string{ + "runtime config:", + "context:", + "project & environment:", + "credentials & connectivity:", + "telemetry:", + "no telemetry yet", + } { + if !strings.Contains(output, want) { + t.Errorf("fresh-session /status: expected output to contain %q", want) + } + } +} + +// TestStatusReflectsModelSwitch verifies /status reads the live run config each +// invocation, so a /model switch (which mutates live.Model/providerName) is +// reflected on the next /status without a restart. +func TestStatusReflectsModelSwitch(t *testing.T) { + host := newFakeHost() + host.live.Model = "model-a" + host.live.ProviderName = "prov-a" + + var buf bytes.Buffer + RunStatus(&buf, host) + if out := buf.String(); !strings.Contains(out, "model: model-a") || !strings.Contains(out, "provider: prov-a") { + t.Errorf("expected model-a/prov-a, got:\n%s", out) + } + + // Simulate a /model switch mutating the live config. + host.live.Model = "model-b" + host.live.ProviderName = "prov-b" + buf.Reset() + RunStatus(&buf, host) + out := buf.String() + if !strings.Contains(out, "model: model-b") || !strings.Contains(out, "provider: prov-b") { + t.Errorf("expected model-b/prov-b after switch, got:\n%s", out) + } + if strings.Contains(out, "model: model-a") { + t.Errorf("stale model-a still present after switch:\n%s", out) + } +} + +// TestStatusTelemetryResetOnFork verifies that after the telemetry holder is +// reset (as runForkClone/runImport do on /fork, /clone, /import - wired in +// #291), /status shows "no telemetry yet" again, so cumulative stats do not +// bleed across conversations. +func TestStatusTelemetryResetOnFork(t *testing.T) { + host := newFakeHost() + + holder := cli.NewTelemetryHolder() + holder.Fold(agentcore.TelemetryEvent{ + Turns: 3, + TruncationCount: 1, + CompactionCount: 0, + ContextUtilization: 0.42, + ContextTokens: 53760, + ContextWindow: 128000, + ToolDurationsMs: map[string]agentcore.ToolTiming{"bash": {Count: 2, TotalMs: 150}}, + }) + host.telemetry = holder + + var buf bytes.Buffer + RunStatus(&buf, host) + if out := buf.String(); !strings.Contains(out, "turns: 3") { + t.Errorf("expected 'turns: 3' before reset, got:\n%s", out) + } + + // /fork, /clone, /import all call holder.Reset() (runForkClone/runImport). + holder.Reset() + buf.Reset() + RunStatus(&buf, host) + out := buf.String() + if n := strings.Count(out, "no telemetry yet"); n != 2 { + t.Errorf("expected 2 'no telemetry yet' after reset (cumulative + last run), got %d:\n%s", n, out) + } + if strings.Contains(out, "turns: 3") { + t.Errorf("stale telemetry after reset:\n%s", out) + } +} + +// TestStatusTiming verifies /status renders fast (target <50ms; assert <100ms +// to absorb CI runner variance). It is pure in-memory rendering - no disk or +// network I/O on the hot path. +func TestStatusTiming(t *testing.T) { + host := newFakeHost() + host.cwd = "/tmp/e2e-timing" + host.live.ContextWindow = 128000 + host.agentCtx.Messages = append(host.agentCtx.Messages, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hello")}}, + ) + holder := cli.NewTelemetryHolder() + holder.Fold(agentcore.TelemetryEvent{ + Turns: 5, + ContextUtilization: 0.5, + ContextTokens: 64000, + ContextWindow: 128000, + ToolDurationsMs: map[string]agentcore.ToolTiming{"bash": {Count: 3, TotalMs: 210}, "read": {Count: 6, TotalMs: 90}}, + }) + host.telemetry = holder + + // Warm once (allocs), then measure. + var buf bytes.Buffer + RunStatus(&buf, host) + buf.Reset() + + start := time.Now() + RunStatus(&buf, host) + elapsed := time.Since(start) + if elapsed >= 100*time.Millisecond { + t.Errorf("/status render took %v, want <100ms (target <50ms)", elapsed) + } +} diff --git a/pigo/internal/cli/status/status_test.go b/pigo/internal/cli/status/status_test.go new file mode 100644 index 0000000..43f9ab2 --- /dev/null +++ b/pigo/internal/cli/status/status_test.go @@ -0,0 +1,254 @@ +package status + +import ( + "bytes" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/runtime" +) + +func TestRunStatus(t *testing.T) { + host := newFakeHost() + host.live.Model = "test-model" + host.live.ProviderName = "test-provider" + host.live.BaseURL = "https://api.example.com" + host.live.Protocol = "anthropic" + host.live.ContextWindow = 128000 + + host.agentCtx.Messages = append(host.agentCtx.Messages, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hello")}}, + ) + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + if !strings.Contains(output, "runtime config:") { + t.Error("expected output to contain 'runtime config:'") + } + if !strings.Contains(output, "model: test-model") { + t.Error("expected output to contain 'model: test-model'") + } + if !strings.Contains(output, "provider: test-provider") { + t.Error("expected output to contain 'provider: test-provider'") + } + if !strings.Contains(output, "base URL: https://api.example.com") { + t.Error("expected output to contain 'base URL: https://api.example.com'") + } + if !strings.Contains(output, "protocol: anthropic") { + t.Error("expected output to contain 'protocol: anthropic'") + } + if !strings.Contains(output, "context window: 128000 tokens") { + t.Error("expected output to contain 'context window: 128000 tokens'") + } + + if !strings.Contains(output, "context:") { + t.Error("expected output to contain 'context:'") + } + if !strings.Contains(output, "current:") { + t.Error("expected output to contain 'current:'") + } + if !strings.Contains(output, "utilization:") { + t.Error("expected output to contain 'utilization:'") + } + if !strings.Contains(output, "compactions: 0") { + t.Error("expected output to contain 'compactions: 0'") + } + if !strings.Contains(output, "before compact:") { + t.Error("expected output to contain 'before compact:'") + } +} + +func TestRunStatusWithCompaction(t *testing.T) { + host := newFakeHost() + host.live.ContextWindow = 128000 + + host.agentCtx.Messages = append(host.agentCtx.Messages, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hello")}}, + agentcore.CompactionMessage{Summary: "compacted history"}, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("more")}}, + ) + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + if !strings.Contains(output, "compactions: 1") { + t.Error("expected output to contain 'compactions: 1'") + } +} + +func TestRunStatusUnknownContextWindow(t *testing.T) { + host := newFakeHost() + host.live.ContextWindow = 0 // unknown + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + if !strings.Contains(output, "context window: unknown") { + t.Error("expected output to contain 'context window: unknown'") + } + if !strings.Contains(output, "auto-compaction disabled") { + t.Error("expected output to contain 'auto-compaction disabled'") + } +} + +func TestBeforeCompactCalculation(t *testing.T) { + reserve := compaction.DefaultCompactionSettings.ReserveTokens + if reserve != 16384 { + t.Errorf("expected reserve tokens to be 16384, got %d", reserve) + } + + contextWindow := 128000 + threshold := contextWindow - reserve + if threshold != 128000-16384 { + t.Errorf("expected threshold to be 128000-16384=%d, got %d", 128000-16384, threshold) + } +} + +func TestRunStatusEnvAndCreds(t *testing.T) { + host := newFakeHost() + host.cwd = "/tmp/test-cwd" + host.live.ProviderName = "test-provider" + host.live.BaseURL = "https://api.example.com" + + host.slash.AddSkill(runtime.SlashCommand{Name: "my-skill", Expand: func(string) string { return "" }}) + host.slash.AddPlugin(runtime.SlashCommand{Name: "my-plugin", Run: func(string) (string, string) { return "", "" }}) + + host.creds.SetOverride("test-provider", "sk-secretkey-wxyz") + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + if !strings.Contains(output, "project & environment:") { + t.Error("expected 'project & environment:' section") + } + if !strings.Contains(output, "cwd: /tmp/test-cwd") { + t.Error("expected 'cwd: /tmp/test-cwd'") + } + if !strings.Contains(output, "trust: disabled") { + t.Error("expected 'trust: disabled' when trust manager is nil") + } + if !strings.Contains(output, "skills: 1 (my-skill)") { + t.Error("expected 'skills: 1 (my-skill)'") + } + if !strings.Contains(output, "plugins: 1 (my-plugin)") { + t.Error("expected 'plugins: 1 (my-plugin)'") + } + + if !strings.Contains(output, "credentials & connectivity:") { + t.Error("expected 'credentials & connectivity:' section") + } + if !strings.Contains(output, "api key: set") { + t.Error("expected 'api key: set'") + } + if !strings.Contains(output, "••••wxyz") { + t.Error("expected masked key '••••wxyz'") + } + if strings.Contains(output, "sk-secretkey-wxyz") { + t.Error("full API key leaked into /status output") + } + if !strings.Contains(output, "endpoint: https://api.example.com") { + t.Error("expected 'endpoint: https://api.example.com'") + } +} + +func TestRunStatusTelemetryNoData(t *testing.T) { + host := newFakeHost() + // host.telemetry is nil. + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + if !strings.Contains(output, "telemetry:") { + t.Error("expected 'telemetry:' section") + } + if n := strings.Count(output, "no telemetry yet"); n != 2 { + t.Errorf("expected 2 'no telemetry yet' (cumulative + last run), got %d", n) + } +} + +func TestRunStatusTelemetryPopulated(t *testing.T) { + host := newFakeHost() + host.live.ContextWindow = 128000 + + holder := cli.NewTelemetryHolder() + holder.Fold(agentcore.TelemetryEvent{ + Turns: 3, + TruncationCount: 1, + CompactionCount: 0, + ContextUtilization: 0.42, + ContextTokens: 53760, + ContextWindow: 128000, + ToolDurationsMs: map[string]agentcore.ToolTiming{ + "bash": {Count: 2, TotalMs: 150}, + "read": {Count: 4, TotalMs: 80}, + }, + }) + host.telemetry = holder + + var buf bytes.Buffer + RunStatus(&buf, host) + output := buf.String() + + if strings.Contains(output, "no telemetry yet") { + t.Error("did not expect 'no telemetry yet' when telemetry is populated") + } + if !strings.Contains(output, "since session start:") { + t.Error("expected 'since session start:' cumulative block") + } + if !strings.Contains(output, "last run:") { + t.Error("expected 'last run:' block") + } + if !strings.Contains(output, "turns: 3") { + t.Error("expected 'turns: 3' (last run == cumulative after one run)") + } + if !strings.Contains(output, "bash") || !strings.Contains(output, "2 calls") || !strings.Contains(output, "150ms") { + t.Error("expected bash tool row with '2 calls' / '150ms'") + } + if !strings.Contains(output, "utilization: 42%") { + t.Error("expected 'utilization: 42%'") + } + + holder.Fold(agentcore.TelemetryEvent{ + Turns: 2, + TruncationCount: 0, + CompactionCount: 1, + ContextUtilization: 0.5, + ContextTokens: 64000, + ContextWindow: 128000, + ToolDurationsMs: map[string]agentcore.ToolTiming{ + "bash": {Count: 1, TotalMs: 40}, + }, + }) + buf.Reset() + RunStatus(&buf, host) + output = buf.String() + if !strings.Contains(output, "turns: 5") { + t.Error("expected cumulative 'turns: 5' after two runs") + } + if !strings.Contains(output, "turns: 2") { + t.Error("expected last-run 'turns: 2' after two runs") + } +} + +func TestMaskKey(t *testing.T) { + cases := []struct{ in, want string }{ + {"sk-secretkey-wxyz", "••••wxyz"}, + {"abcd", "••••"}, // exactly 4 -> masked entirely + {"ab", "••"}, + {"", ""}, + } + for _, c := range cases { + if got := maskKey(c.in); got != c.want { + t.Errorf("maskKey(%q) = %q, want %q", c.in, got, c.want) + } + } +} diff --git a/pigo/internal/cli/telemetry.go b/pigo/internal/cli/telemetry.go new file mode 100644 index 0000000..9330120 --- /dev/null +++ b/pigo/internal/cli/telemetry.go @@ -0,0 +1,135 @@ +// This file implements the telemetry holder that retains per-run telemetry +// events (US-001, #291) and folds them into a cumulative accumulator for the +// /status command to display. It was moved verbatim from cmd/pigo +// (telemetry_state.go, US-013) into the shared cli layer so the status and repl +// subpackages share one holder type. +package cli + +import ( + "github.com/smallnest/pigo/internal/agentcore" +) + +// TelemetryHolder holds the most recent run's telemetry event and a cumulative +// accumulator that sums metrics across all runs in the session. +type TelemetryHolder struct { + // last is the telemetry event from the most recent completed run, or nil if + // no run has completed yet. + last *agentcore.TelemetryEvent + // cumulative is the accumulator that sums metrics across all runs. + cumulative cumulativeTelemetry +} + +// cumulativeTelemetry holds the accumulated metrics from all runs in the session. +type cumulativeTelemetry struct { + // turns is the total number of turns across all runs. + turns int + // toolDurationsMs maps tool name to aggregated count and total milliseconds + // across all runs. + toolDurationsMs map[string]agentcore.ToolTiming + // truncationCount is the total number of truncations across all runs. + truncationCount int + // compactionCount is the total number of compactions across all runs. + compactionCount int + // contextTokens is the most recent observed context token count. + contextTokens int + // contextWindow is the most recent observed context window size. + contextWindow int +} + +// NewTelemetryHolder creates a new, empty TelemetryHolder with no telemetry yet. +func NewTelemetryHolder() *TelemetryHolder { + return &TelemetryHolder{ + last: nil, + cumulative: cumulativeTelemetry{ + toolDurationsMs: make(map[string]agentcore.ToolTiming), + }, + } +} + +// Reset clears the holder to "no telemetry yet", as if the session had just +// started. It is called when switching sessions (/fork, /clone, /import). +func (h *TelemetryHolder) Reset() { + h.last = nil + h.cumulative = cumulativeTelemetry{ + toolDurationsMs: make(map[string]agentcore.ToolTiming), + } +} + +// Fold incorporates a new TelemetryEvent into the holder: it becomes the new +// "last run" event, and its metrics are added to the cumulative accumulator. +func (h *TelemetryHolder) Fold(ev agentcore.TelemetryEvent) { + // Store the new event as the last run. + h.last = &ev + + // Add its metrics to the cumulative accumulator. + h.cumulative.turns += ev.Turns + h.cumulative.truncationCount += ev.TruncationCount + h.cumulative.compactionCount += ev.CompactionCount + + // Update the latest context tokens and window (always use the most recent). + h.cumulative.contextTokens = ev.ContextTokens + h.cumulative.contextWindow = ev.ContextWindow + + // Sum the per-tool timings. + for name, timing := range ev.ToolDurationsMs { + agg := h.cumulative.toolDurationsMs[name] + agg.Count += timing.Count + agg.TotalMs += timing.TotalMs + h.cumulative.toolDurationsMs[name] = agg + } +} + +// HasTelemetry returns true if at least one run has completed and telemetry is +// available. +func (h *TelemetryHolder) HasTelemetry() bool { + return h.last != nil +} + +// Last returns the most recent TelemetryEvent, or nil if no run has completed. +func (h *TelemetryHolder) Last() *agentcore.TelemetryEvent { + return h.last +} + +// CumulativeTurns returns the total number of turns across all runs. +func (h *TelemetryHolder) CumulativeTurns() int { + return h.cumulative.turns +} + +// CumulativeTruncationCount returns the total number of truncations across all runs. +func (h *TelemetryHolder) CumulativeTruncationCount() int { + return h.cumulative.truncationCount +} + +// CumulativeCompactionCount returns the total number of compactions across all runs. +func (h *TelemetryHolder) CumulativeCompactionCount() int { + return h.cumulative.compactionCount +} + +// CumulativeToolDurations returns the aggregated tool timings across all runs. +func (h *TelemetryHolder) CumulativeToolDurations() map[string]agentcore.ToolTiming { + // Return a copy to prevent mutation of the internal state. + result := make(map[string]agentcore.ToolTiming, len(h.cumulative.toolDurationsMs)) + for name, timing := range h.cumulative.toolDurationsMs { + result[name] = timing + } + return result +} + +// LatestContextTokens returns the most recent observed context token count. +func (h *TelemetryHolder) LatestContextTokens() int { + return h.cumulative.contextTokens +} + +// LatestContextWindow returns the most recent observed context window size. +func (h *TelemetryHolder) LatestContextWindow() int { + return h.cumulative.contextWindow +} + +// CumulativeContextUtilization returns the latest utilization ratio, or 0 if no +// window is known. +func (h *TelemetryHolder) CumulativeContextUtilization() float64 { + if h.cumulative.contextWindow <= 0 { + return 0 + } + return float64(h.cumulative.contextTokens) / float64(h.cumulative.contextWindow) +} diff --git a/pigo/internal/cli/telemetry_test.go b/pigo/internal/cli/telemetry_test.go new file mode 100644 index 0000000..d6f9a93 --- /dev/null +++ b/pigo/internal/cli/telemetry_test.go @@ -0,0 +1,209 @@ +package cli + +import ( + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// TestTelemetryHolderReset verifies that Reset clears the holder to "no telemetry yet". +func TestTelemetryHolderReset(t *testing.T) { + h := NewTelemetryHolder() + if h.HasTelemetry() { + t.Fatal("expected HasTelemetry() = false on new holder") + } + + // Fold an event and verify it's present. + h.Fold(agentcore.TelemetryEvent{ + Turns: 1, + TruncationCount: 1, + CompactionCount: 1, + ContextTokens: 100, + ContextWindow: 1000, + }) + if !h.HasTelemetry() { + t.Fatal("expected HasTelemetry() = true after Fold") + } + if h.Last() == nil { + t.Fatal("expected Last() != nil after Fold") + } + if h.CumulativeTurns() != 1 { + t.Fatalf("expected CumulativeTurns() = 1, got %d", h.CumulativeTurns()) + } + + // Reset and verify it's cleared. + h.Reset() + if h.HasTelemetry() { + t.Fatal("expected HasTelemetry() = false after Reset") + } + if h.Last() != nil { + t.Fatal("expected Last() = nil after Reset") + } + if h.CumulativeTurns() != 0 { + t.Fatalf("expected CumulativeTurns() = 0 after Reset, got %d", h.CumulativeTurns()) + } + if h.CumulativeTruncationCount() != 0 { + t.Fatalf("expected CumulativeTruncationCount() = 0 after Reset, got %d", h.CumulativeTruncationCount()) + } + if h.CumulativeCompactionCount() != 0 { + t.Fatalf("expected CumulativeCompactionCount() = 0 after Reset, got %d", h.CumulativeCompactionCount()) + } + if h.LatestContextTokens() != 0 { + t.Fatalf("expected LatestContextTokens() = 0 after Reset, got %d", h.LatestContextTokens()) + } + if h.LatestContextWindow() != 0 { + t.Fatalf("expected LatestContextWindow() = 0 after Reset, got %d", h.LatestContextWindow()) + } +} + +// TestTelemetryHolderFoldTwoEvents verifies that folding two synthetic TelemetryEvents +// sums the metrics correctly. +func TestTelemetryHolderFoldTwoEvents(t *testing.T) { + h := NewTelemetryHolder() + + // Fold first event. + h.Fold(agentcore.TelemetryEvent{ + Turns: 2, + TruncationCount: 1, + CompactionCount: 3, + ContextTokens: 100, + ContextWindow: 1000, + ToolDurationsMs: map[string]agentcore.ToolTiming{ + "foo": {Count: 1, TotalMs: 100}, + "bar": {Count: 2, TotalMs: 200}, + }, + }) + + // Verify the first event is the last one. + if !h.HasTelemetry() { + t.Fatal("expected HasTelemetry() = true after first Fold") + } + last1 := h.Last() + if last1 == nil { + t.Fatal("expected Last() != nil after first Fold") + } + if last1.Turns != 2 { + t.Fatalf("expected Last().Turns = 2, got %d", last1.Turns) + } + + // Verify cumulative after first event. + if h.CumulativeTurns() != 2 { + t.Fatalf("expected CumulativeTurns() = 2 after first Fold, got %d", h.CumulativeTurns()) + } + if h.CumulativeTruncationCount() != 1 { + t.Fatalf("expected CumulativeTruncationCount() = 1 after first Fold, got %d", h.CumulativeTruncationCount()) + } + if h.CumulativeCompactionCount() != 3 { + t.Fatalf("expected CumulativeCompactionCount() = 3 after first Fold, got %d", h.CumulativeCompactionCount()) + } + if h.LatestContextTokens() != 100 { + t.Fatalf("expected LatestContextTokens() = 100 after first Fold, got %d", h.LatestContextTokens()) + } + if h.LatestContextWindow() != 1000 { + t.Fatalf("expected LatestContextWindow() = 1000 after first Fold, got %d", h.LatestContextWindow()) + } + tools1 := h.CumulativeToolDurations() + if tools1["foo"].Count != 1 || tools1["foo"].TotalMs != 100 { + t.Fatalf("expected foo tool timing {1, 100}, got %+v", tools1["foo"]) + } + if tools1["bar"].Count != 2 || tools1["bar"].TotalMs != 200 { + t.Fatalf("expected bar tool timing {2, 200}, got %+v", tools1["bar"]) + } + + // Fold second event. + h.Fold(agentcore.TelemetryEvent{ + Turns: 3, + TruncationCount: 2, + CompactionCount: 1, + ContextTokens: 150, + ContextWindow: 1000, + ToolDurationsMs: map[string]agentcore.ToolTiming{ + "foo": {Count: 2, TotalMs: 250}, + "baz": {Count: 1, TotalMs: 50}, + }, + }) + + // Verify the second event is now the last one. + last2 := h.Last() + if last2 == nil { + t.Fatal("expected Last() != nil after second Fold") + } + if last2.Turns != 3 { + t.Fatalf("expected Last().Turns = 3, got %d", last2.Turns) + } + + // Verify cumulative after second event (sums of both). + if h.CumulativeTurns() != 5 { + t.Fatalf("expected CumulativeTurns() = 5 after second Fold, got %d", h.CumulativeTurns()) + } + if h.CumulativeTruncationCount() != 3 { + t.Fatalf("expected CumulativeTruncationCount() = 3 after second Fold, got %d", h.CumulativeTruncationCount()) + } + if h.CumulativeCompactionCount() != 4 { + t.Fatalf("expected CumulativeCompactionCount() = 4 after second Fold, got %d", h.CumulativeCompactionCount()) + } + if h.LatestContextTokens() != 150 { + t.Fatalf("expected LatestContextTokens() = 150 after second Fold, got %d", h.LatestContextTokens()) + } + if h.LatestContextWindow() != 1000 { + t.Fatalf("expected LatestContextWindow() = 1000 after second Fold, got %d", h.LatestContextWindow()) + } + tools2 := h.CumulativeToolDurations() + if tools2["foo"].Count != 3 || tools2["foo"].TotalMs != 350 { + t.Fatalf("expected foo tool timing {3, 350} after sum, got %+v", tools2["foo"]) + } + if tools2["bar"].Count != 2 || tools2["bar"].TotalMs != 200 { + t.Fatalf("expected bar tool timing {2, 200} unchanged, got %+v", tools2["bar"]) + } + if tools2["baz"].Count != 1 || tools2["baz"].TotalMs != 50 { + t.Fatalf("expected baz tool timing {1, 50} after sum, got %+v", tools2["baz"]) + } +} + +// TestTelemetryHolderCumulativeContextUtilization verifies context utilization calculation. +func TestTelemetryHolderCumulativeContextUtilization(t *testing.T) { + h := NewTelemetryHolder() + + // With unknown window (0), utilization should be 0. + h.Fold(agentcore.TelemetryEvent{ + ContextTokens: 100, + ContextWindow: 0, + }) + if h.CumulativeContextUtilization() != 0 { + t.Fatalf("expected utilization = 0 with unknown window, got %f", h.CumulativeContextUtilization()) + } + + // With known window, utilization should be tokens/window. + h.Fold(agentcore.TelemetryEvent{ + ContextTokens: 500, + ContextWindow: 1000, + }) + if h.CumulativeContextUtilization() != 0.5 { + t.Fatalf("expected utilization = 0.5, got %f", h.CumulativeContextUtilization()) + } +} + +// TestTelemetryHolderToolDurationsImmutable verifies that CumulativeToolDurations returns +// a copy to prevent external mutation. +func TestTelemetryHolderToolDurationsImmutable(t *testing.T) { + h := NewTelemetryHolder() + h.Fold(agentcore.TelemetryEvent{ + ToolDurationsMs: map[string]agentcore.ToolTiming{ + "test": {Count: 1, TotalMs: 100}, + }, + }) + + // Get the tool durations and mutate the returned map. + tools := h.CumulativeToolDurations() + tools["test"] = agentcore.ToolTiming{Count: 999, TotalMs: 9999} + tools["new"] = agentcore.ToolTiming{Count: 1, TotalMs: 1} + + // Verify the internal state was not modified. + tools2 := h.CumulativeToolDurations() + if tools2["test"].Count != 1 || tools2["test"].TotalMs != 100 { + t.Fatalf("expected internal tool timing to remain {1, 100}, got %+v", tools2["test"]) + } + if _, ok := tools2["new"]; ok { + t.Fatal("expected 'new' tool not to be present in internal state") + } +} diff --git a/pigo/internal/cli/testutil/prompts.go b/pigo/internal/cli/testutil/prompts.go new file mode 100644 index 0000000..82d3986 --- /dev/null +++ b/pigo/internal/cli/testutil/prompts.go @@ -0,0 +1,24 @@ +// Package testutil holds test helpers shared across the internal/cli +// subpackages. Only helpers reused by two or more test files and free of +// cmd/pigo-internal types belong here; helpers that build the concrete replDeps +// aggregate stay with the package that owns it. +package testutil + +import ( + "os" + "path/filepath" + "testing" +) + +// WritePrompt writes a prompt-template .md at home// with the given +// body, creating the directory as needed. It fails the test on any I/O error. +func WritePrompt(t *testing.T, home, sub, name, body string) { + t.Helper() + d := filepath.Join(home, sub) + if err := os.MkdirAll(d, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(d, name), []byte(body), 0o644); err != nil { + t.Fatal(err) + } +} diff --git a/pigo/internal/cli/tui/banner.go b/pigo/internal/cli/tui/banner.go new file mode 100644 index 0000000..852555a --- /dev/null +++ b/pigo/internal/cli/tui/banner.go @@ -0,0 +1,105 @@ +package tui + +import ( + "fmt" + "strings" + + "charm.land/lipgloss/v2" + + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/selfupdate" +) + +// This file builds the startup splash shown at the top of the transcript: the +// pigo braille logo painted in a vertical rainbow gradient, with the session's +// basic configuration (model, provider, protocol, thinking effort, directory) +// laid out beside it. It is seeded once by withSession so it scrolls up as the +// conversation grows, like a shell's login banner. + +// logoLines is the pigo braille-art logo, one string per row. +var logoLines = []string{ + "⣿⣿⣿⣿⡿⠟⠛⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠙⣿", + "⣿⣿⡿⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣼", + "⣿⠋⠀⠀⠀⠀⠀⣀⡀⠀⠀⠀⠀⢠⣤⣤⣤⠀⠀⠀⠀⠀⣤⣤⣤⣤⣤⣤⣾⣿", + "⣧⠀⠀⠀⠀⣠⣾⣿⡇⠀⠀⠀⠀⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⣿⣿⣿⣿⣿⣿", + "⣿⣶⣤⣤⣾⣿⣿⣿⡇⠀⠀⠀⠀⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⣿⣿⣿⣿⣿⣿", + "⣿⣿⣿⣿⣿⣿⣿⣿⠀⠀⠀⠀⢀⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⣿⣿⣿⣿⣿⣿", + "⣿⣿⣿⣿⣿⣿⣿⡟⠀⠀⠀⠀⢸⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⣿⣿⣿⣿⣿⣿", + "⣿⣿⣿⣿⣿⣿⣿⠇⠀⠀⠀⠀⣾⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⣿⣿⣿⣿⣿⣿", + "⣿⣿⣿⣿⣿⣿⡟⠀⠀⠀⠀⢠⣿⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⣿⡿⠛⠛⢿⣿", + "⣿⣿⣿⣿⣿⡿⠁⠀⠀⠀⠀⣾⣿⣿⣿⣿⣿⠀⠀⠀⠀⠀⣿⣿⡟⠀⠀⠀⠀⢻", + "⣿⣿⣿⣿⡟⠁⠀⠀⠀⠀⣸⣿⣿⣿⣿⣿⣿⡀⠀⠀⠀⠀⠛⠛⠁⠀⠀⠀⠀⣾", + "⣿⣿⣿⡏⠀⠀⠀⠀⠀⣴⣿⣿⣿⣿⣿⣿⣿⣧⡀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣼⣿", + "⣿⣿⣿⣿⣄⣀⣀⣠⣾⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣦⣄⣀⣀⣀⣀⣠⣴⣿⣿⣿", +} + +// logoColors is the top-to-bottom rainbow ramp painted across the logo rows +// (ANSI 256-color cube): red → orange → yellow → green → cyan → blue. +var logoColors = []string{ + "196", "202", "208", "214", "220", "190", "118", + "46", "48", "50", "45", "39", "33", +} + +// renderBanner paints the logo gradient and joins it with a config panel showing +// the session basics. Its only I/O is a single cheap read of the local +// update-check cache (no network — CachedLatest); it never panics, so it is safe +// to build eagerly at startup. +func renderBanner(theme Theme, opts Options, cwd string) string { + var logo strings.Builder + for i, line := range logoLines { + if i > 0 { + logo.WriteByte('\n') + } + c := logoColors[i%len(logoColors)] + logo.WriteString(lipgloss.NewStyle().Foreground(lipgloss.Color(c)).Render(line)) + } + + title := lipgloss.NewStyle().Foreground(lipgloss.Color(colorAccent)).Bold(true) + label := lipgloss.NewStyle().Foreground(lipgloss.Color(colorGray)) + value := lipgloss.NewStyle().Foreground(lipgloss.Color(colorUser)).Bold(true) + + rows := [][2]string{ + {"Version", firstNonEmpty(opts.Version, "dev")}, + {"Model", firstNonEmpty(opts.Model, "—")}, + {"Provider", firstNonEmpty(opts.ProviderName, "—")}, + {"Protocol", firstNonEmpty(provider.ProtocolLabel(opts.Protocol), "—")}, + {"Thinking", firstNonEmpty(string(opts.ThinkingLevel), "off")}, + {"Directory", firstNonEmpty(cwd, "—")}, + } + + // When the cached latest-release check says a newer version exists, append a + // highlighted "→ vX.Y.Z" and an upgrade hint to the Version row. The check is + // read from the local cache only (no network here); a background refresh keeps + // it current for the next launch. dev/unparseable versions never trigger this. + upgradeHint := "" + if latest, _ := selfupdate.CachedLatest(); latest != "" { + if avail, comparable := selfupdate.UpdateAvailable(opts.Version, latest); comparable && avail { + newVer := lipgloss.NewStyle().Foreground(lipgloss.Color("214")).Bold(true).Render(latest) + rows[0][1] = rows[0][1] + " → " + newVer + upgradeHint = label.Render(strings.Repeat(" ", 11)) + + lipgloss.NewStyle().Foreground(lipgloss.Color("214")).Render("Run pigo update to upgrade") + } + } + + var info strings.Builder + info.WriteString(title.Render("pigo") + " " + theme.System.Render("Terminal AI coding assistant") + "\n\n") + for i, r := range rows { + if i > 0 { + info.WriteByte('\n') + } + info.WriteString(label.Render(fmt.Sprintf("%-10s ", r[0])) + value.Render(r[1])) + } + if upgradeHint != "" { + info.WriteString("\n" + upgradeHint) + } + + return lipgloss.JoinHorizontal(lipgloss.Center, logo.String(), " ", info.String()) +} + +// firstNonEmpty returns s when it is non-empty, otherwise the fallback. +func firstNonEmpty(s, fallback string) string { + if strings.TrimSpace(s) == "" { + return fallback + } + return s +} diff --git a/pigo/internal/cli/tui/banner_test.go b/pigo/internal/cli/tui/banner_test.go new file mode 100644 index 0000000..2734ac1 --- /dev/null +++ b/pigo/internal/cli/tui/banner_test.go @@ -0,0 +1,90 @@ +package tui + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// writeUpdateCache seeds the selfupdate cache under a temp PIGO_HOME so the +// banner's cached-latest lookup is deterministic. +func writeUpdateCache(t *testing.T, latest string) { + t.Helper() + dir := t.TempDir() + t.Setenv("PIGO_HOME", dir) + data, _ := json.Marshal(map[string]any{ + "checked_at": time.Now(), + "latest": latest, + }) + if err := os.WriteFile(filepath.Join(dir, "update-check.json"), data, 0o644); err != nil { + t.Fatal(err) + } +} + +func TestRenderBannerShowsVersion(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) // empty cache: no upgrade hint + out := renderBanner(DefaultTheme(), Options{Version: "v0.3.1"}, "/tmp/proj") + if !strings.Contains(out, "Version") || !strings.Contains(out, "v0.3.1") { + t.Errorf("banner missing Version row: %q", out) + } + if strings.Contains(out, "Run pigo update to upgrade") { + t.Error("banner should not show upgrade hint with empty cache") + } +} + +func TestRenderBannerDevNoHint(t *testing.T) { + writeUpdateCache(t, "v9.9.9") // even with a newer tag cached... + out := renderBanner(DefaultTheme(), Options{Version: "dev"}, "/tmp/proj") + if !strings.Contains(out, "dev") { + t.Errorf("banner should show dev version: %q", out) + } + if strings.Contains(out, "Run pigo update to upgrade") { + t.Error("dev build must not show an upgrade hint") + } +} + +func TestRenderBannerUpgradeHint(t *testing.T) { + writeUpdateCache(t, "v0.4.0") + out := renderBanner(DefaultTheme(), Options{Version: "v0.3.1"}, "/tmp/proj") + if !strings.Contains(out, "v0.4.0") { + t.Errorf("banner should highlight newer version v0.4.0: %q", out) + } + if !strings.Contains(out, "Run pigo update to upgrade") { + t.Errorf("banner should show upgrade hint: %q", out) + } +} + +func TestRenderBannerUpToDate(t *testing.T) { + writeUpdateCache(t, "v0.3.1") + out := renderBanner(DefaultTheme(), Options{Version: "v0.3.1"}, "/tmp/proj") + if strings.Contains(out, "Run pigo update to upgrade") { + t.Error("up-to-date build must not show an upgrade hint") + } +} + +// TestRenderBannerProtocolLabel verifies the Protocol row shows the concrete +// OpenAI wire variant: a bare "openai" is surfaced as "openai/chat" (explicit +// Chat Completions), "openai/resp_api" passes through, and an unset protocol +// falls back to the em dash rather than showing an empty row. +func TestRenderBannerProtocolLabel(t *testing.T) { + cases := []struct { + protocol string + want string + }{ + {"openai", "openai/chat"}, + {"openai/chat", "openai/chat"}, + {"openai/resp_api", "openai/resp_api"}, + {"anthropic", "anthropic"}, + {"", "—"}, + } + for _, c := range cases { + t.Setenv("PIGO_HOME", t.TempDir()) + out := renderBanner(DefaultTheme(), Options{Version: "dev", Protocol: c.protocol}, "/tmp/proj") + if !strings.Contains(out, c.want) { + t.Errorf("protocol %q: banner should show %q, got: %q", c.protocol, c.want, out) + } + } +} diff --git a/pigo/internal/cli/tui/bridge.go b/pigo/internal/cli/tui/bridge.go new file mode 100644 index 0000000..8e8c64f --- /dev/null +++ b/pigo/internal/cli/tui/bridge.go @@ -0,0 +1,139 @@ +package tui + +import ( + "context" + "encoding/json" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/runtime" +) + +// This file bridges the agent run seam (runtime.StartRun + runtime.DrainStream) +// to Bubble Tea (US-004, SPEC 5.1 bridge / 3.2). The agent loop runs on its own +// goroutine and emits AgentEvents; a Bubble Tea program consumes tea.Msg values +// one at a time from its Update loop. The bridge is a pump: a goroutine drains +// the run and converts every event into the matching tea.Msg (see msgs.go), +// sending it into a buffered channel; a tea.Cmd (waitForEvent) receives one msg +// per Update tick. The channel is the only synchronization point, so the +// producer never touches the model and the model never touches the run — all +// state transitions happen on the tea goroutine. +// +// Back-pressure is intentional: the channel blocks the draining goroutine when +// the buffer is full, so no event is ever dropped (the tea loop always catches +// up). Node #388 wires startRun into Model.Init/Update; this file only provides +// the reusable, unit-testable primitives. + +// eventChanCap is the buffer size of the bridge channel. A modest buffer lets a +// burst of tool events queue without blocking the run's goroutine on every send, +// while still bounding memory (blocking, never dropping, past the cap). +const eventChanCap = 64 + +// newEventChan allocates the buffered channel the bridge pumps run events +// through. +func newEventChan() chan tea.Msg { + return make(chan tea.Msg, eventChanCap) +} + +// newStreamHandler builds the runtime.StreamHandler that converts each run event +// into a tea.Msg and sends it into ch. Sends block when ch is full, applying +// back-pressure to the draining goroutine so no event is lost. It is factored +// out of pump so the callback→msg conversion can be unit-tested without a real +// provider run (see bridge_test.go). +func newStreamHandler(ch chan tea.Msg, extra func(agentcore.AgentEvent)) runtime.StreamHandler { + return runtime.StreamHandler{ + OnText: func(delta string) { + ch <- textDeltaMsg{delta: delta} + }, + OnTurnEnd: func(msg agentcore.AssistantMessage, results []agentcore.ToolResultMessage) { + ch <- turnEndMsg{msg: msg, results: results} + }, + OnEvent: func(ev agentcore.AgentEvent) { + // Deliver observer events (plugin notifier, SessionEnd/PreCompact hook) + // first, then translate into TUI messages. + if extra != nil { + extra(ev) + } + switch e := ev.(type) { + case agentcore.ToolExecutionStartEvent: + ch <- toolStartMsg{id: e.ToolCallID, name: e.ToolName, input: argsToMap(e.Args)} + case agentcore.ToolExecutionUpdateEvent: + ch <- toolUpdateMsg{id: e.ToolCallID, partial: agentcore.ContentToText(e.PartialResult.Content)} + case agentcore.ToolExecutionEndEvent: + ch <- toolEndMsg{id: e.ToolCallID, ok: !e.IsError, result: agentcore.ContentToText(e.Result.Content)} + case agentcore.SubAgentProgressEvent: + ch <- subagentProgressMsg{id: e.ToolCallID, desc: e.Description, activity: e.Activity, tokens: e.Tokens} + case agentcore.TelemetryEvent: + ch <- telemetryMsg{ev: e} + case agentcore.CompactionStartEvent: + ch <- compactionStartMsg{} + case agentcore.CompactionEvent: + ch <- compactionMsg{} + } + }, + } +} + +// argsToMap coerces a tool call's untyped Args into a map[string]any. The event +// layer carries Args as an untyped any: the tool executor emits it as a +// json.RawMessage (the raw decoded JSON arguments), but a caller may also hand +// an already-decoded map. Both are supported here so the tool card can show the +// call's arguments; anything that is not a JSON object yields nil. +func argsToMap(args any) map[string]any { + switch v := args.(type) { + case map[string]any: + return v + case json.RawMessage: + return unmarshalArgsMap(v) + case []byte: + return unmarshalArgsMap(v) + case string: + return unmarshalArgsMap([]byte(v)) + } + return nil +} + +// unmarshalArgsMap parses JSON object bytes into a map, returning nil for empty +// input or anything that is not a JSON object. +func unmarshalArgsMap(b []byte) map[string]any { + if len(b) == 0 { + return nil + } + var m map[string]any + if err := json.Unmarshal(b, &m); err != nil { + return nil + } + return m +} + +// pump runs the agent loop to completion on the calling goroutine, converting +// every event to a tea.Msg on ch, and finally sends a runEndMsg carrying the +// run's result error. It is meant to be launched as a goroutine by startRun. +func pump(ctx context.Context, ch chan tea.Msg, agentCtx *agentcore.AgentContext, cfg runtime.RunConfig, onEvent func(agentcore.AgentEvent)) { + stream := runtime.StartRun(ctx, agentCtx, cfg) + _, err := runtime.DrainStream(ctx, stream, newStreamHandler(ch, onEvent)) + ch <- runEndMsg{err: err} +} + +// waitForEvent returns a tea.Cmd that blocks until the next bridge msg arrives. +// The Update loop re-issues it after handling each msg (except runEndMsg) to +// keep pulling events one at a time, so ordering is preserved and the tea +// goroutine never spins. +func waitForEvent(ch chan tea.Msg) tea.Cmd { + return func() tea.Msg { + return <-ch + } +} + +// startRun launches the run pump on a new goroutine and returns the channel it +// feeds together with the first waitForEvent Cmd. The caller (node #388's model) +// stores the channel and, on every subsequent event, issues waitForEvent(ch) +// again to pull the next msg. Returning the channel keeps the bridge +// self-contained: the model owns the handle and decides when to stop pulling +// (after runEndMsg). +func startRun(ctx context.Context, agentCtx *agentcore.AgentContext, cfg runtime.RunConfig, onEvent func(agentcore.AgentEvent)) (chan tea.Msg, tea.Cmd) { + ch := newEventChan() + go pump(ctx, ch, agentCtx, cfg, onEvent) + return ch, waitForEvent(ch) +} diff --git a/pigo/internal/cli/tui/bridge_test.go b/pigo/internal/cli/tui/bridge_test.go new file mode 100644 index 0000000..795299e --- /dev/null +++ b/pigo/internal/cli/tui/bridge_test.go @@ -0,0 +1,181 @@ +package tui + +import ( + "encoding/json" + "errors" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// drain collects every msg queued on ch without blocking, stopping at the first +// would-block. The bridge sends are synchronous into a buffered channel, so once +// the fake sequence has been driven the msgs are all present and this returns +// them in order. +func drain(ch chan tea.Msg) []tea.Msg { + var out []tea.Msg + for { + select { + case m := <-ch: + out = append(out, m) + default: + return out + } + } +} + +// TestStreamHandlerConversion drives a synthetic event sequence directly through +// the StreamHandler the bridge builds and asserts each callback produces the +// matching tea.Msg, in order. It fakes the run entirely (no provider), exercising +// the callback→msg conversion + channel ordering in isolation. +func TestStreamHandlerConversion(t *testing.T) { + ch := newEventChan() + h := newStreamHandler(ch, nil) + + // A representative sequence: two text deltas, a tool start/update/end, a + // telemetry summary, a compaction, and a turn end. + h.OnText("Hello ") + h.OnText("world") + h.OnEvent(agentcore.ToolExecutionStartEvent{ + ToolCallID: "call-1", + ToolName: "read_file", + Args: map[string]any{"path": "/tmp/x"}, + }) + h.OnEvent(agentcore.ToolExecutionUpdateEvent{ + ToolCallID: "call-1", + ToolName: "read_file", + PartialResult: agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("partial")}}, + }) + h.OnEvent(agentcore.ToolExecutionEndEvent{ + ToolCallID: "call-1", + ToolName: "read_file", + Result: agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("done")}}, + IsError: false, + }) + h.OnEvent(agentcore.TelemetryEvent{Turns: 3}) + h.OnEvent(agentcore.CompactionEvent{Reason: "threshold"}) + h.OnTurnEnd( + agentcore.AssistantMessage{Content: agentcore.ContentList{agentcore.NewTextContent("Hello world")}}, + []agentcore.ToolResultMessage{{ToolCallID: "call-1"}}, + ) + // Simulate drain-done: the pump appends runEndMsg after DrainStream returns. + ch <- runEndMsg{err: nil} + + got := drain(ch) + + if len(got) != 9 { + t.Fatalf("expected 9 msgs, got %d: %#v", len(got), got) + } + + if m, ok := got[0].(textDeltaMsg); !ok || m.delta != "Hello " { + t.Errorf("msg[0] = %#v, want textDeltaMsg{delta:%q}", got[0], "Hello ") + } + if m, ok := got[1].(textDeltaMsg); !ok || m.delta != "world" { + t.Errorf("msg[1] = %#v, want textDeltaMsg{delta:%q}", got[1], "world") + } + if m, ok := got[2].(toolStartMsg); !ok || m.id != "call-1" || m.name != "read_file" || m.input["path"] != "/tmp/x" { + t.Errorf("msg[2] = %#v, want toolStartMsg for call-1", got[2]) + } + if m, ok := got[3].(toolUpdateMsg); !ok || m.id != "call-1" || m.partial != "partial" { + t.Errorf("msg[3] = %#v, want toolUpdateMsg{partial:%q}", got[3], "partial") + } + if m, ok := got[4].(toolEndMsg); !ok || m.id != "call-1" || !m.ok || m.result != "done" { + t.Errorf("msg[4] = %#v, want toolEndMsg{ok:true, result:%q}", got[4], "done") + } + if m, ok := got[5].(telemetryMsg); !ok || m.ev.Turns != 3 { + t.Errorf("msg[5] = %#v, want telemetryMsg{Turns:3}", got[5]) + } + if _, ok := got[6].(compactionMsg); !ok { + t.Errorf("msg[6] = %#v, want compactionMsg", got[6]) + } + if m, ok := got[7].(turnEndMsg); !ok || len(m.results) != 1 || m.results[0].ToolCallID != "call-1" { + t.Errorf("msg[7] = %#v, want turnEndMsg with one result", got[7]) + } + if m, ok := got[8].(runEndMsg); !ok || m.err != nil { + t.Errorf("msg[8] = %#v, want runEndMsg{err:nil}", got[8]) + } +} + +// TestToolEndError verifies the ok flag inverts IsError. +func TestToolEndError(t *testing.T) { + ch := newEventChan() + h := newStreamHandler(ch, nil) + h.OnEvent(agentcore.ToolExecutionEndEvent{ + ToolCallID: "c", + Result: agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("boom")}}, + IsError: true, + }) + got := drain(ch) + if len(got) != 1 { + t.Fatalf("expected 1 msg, got %d", len(got)) + } + m, ok := got[0].(toolEndMsg) + if !ok || m.ok || m.result != "boom" { + t.Errorf("got %#v, want toolEndMsg{ok:false, result:%q}", got[0], "boom") + } +} + +// TestArgsToMap covers the object / non-object coercion of tool call args. +func TestArgsToMap(t *testing.T) { + if m := argsToMap(map[string]any{"k": "v"}); m == nil || m["k"] != "v" { + t.Errorf("object args: got %#v, want map with k=v", m) + } + if m := argsToMap("not-an-object"); m != nil { + t.Errorf("non-object args: got %#v, want nil", m) + } + if m := argsToMap(nil); m != nil { + t.Errorf("nil args: got %#v, want nil", m) + } + // The tool executor emits Args as json.RawMessage; the card must decode it. + if m := argsToMap(json.RawMessage(`{"command":"ls -la"}`)); m == nil || m["command"] != "ls -la" { + t.Errorf("raw JSON object args: got %#v, want map with command=ls -la", m) + } + if m := argsToMap(json.RawMessage(`"just a string"`)); m != nil { + t.Errorf("raw JSON non-object: got %#v, want nil", m) + } + if m := argsToMap(json.RawMessage(nil)); m != nil { + t.Errorf("empty raw JSON: got %#v, want nil", m) + } +} + +// TestSubAgentProgressConversion verifies a SubAgentProgressEvent maps to a +// subagentProgressMsg carrying the id (parent task tool-call id), description, +// activity, and token estimate. +func TestSubAgentProgressConversion(t *testing.T) { + ch := newEventChan() + h := newStreamHandler(ch, nil) + h.OnEvent(agentcore.SubAgentProgressEvent{ + ToolCallID: "task-1", + Description: "build parser", + Activity: "Editing", + Tokens: 256, + }) + got := drain(ch) + if len(got) != 1 { + t.Fatalf("expected 1 msg, got %d", len(got)) + } + m, ok := got[0].(subagentProgressMsg) + if !ok { + t.Fatalf("got %#v, want subagentProgressMsg", got[0]) + } + if m.id != "task-1" || m.desc != "build parser" || m.activity != "Editing" || m.tokens != 256 { + t.Errorf("got %#v, want {id:task-1 desc:build parser activity:Editing tokens:256}", m) + } +} + +// TestWaitForEvent verifies the pump Cmd returns the next queued msg. +func TestWaitForEvent(t *testing.T) { + ch := newEventChan() + want := runEndMsg{err: errors.New("stop")} + ch <- want + cmd := waitForEvent(ch) + if cmd == nil { + t.Fatal("waitForEvent returned nil Cmd") + } + got, ok := cmd().(runEndMsg) + if !ok || got.err == nil || got.err.Error() != "stop" { + t.Errorf("cmd() = %#v, want runEndMsg{err:stop}", got) + } +} diff --git a/pigo/internal/cli/tui/clipimage.go b/pigo/internal/cli/tui/clipimage.go new file mode 100644 index 0000000..2105bac --- /dev/null +++ b/pigo/internal/cli/tui/clipimage.go @@ -0,0 +1,122 @@ +package tui + +import ( + "os" + "os/exec" + "runtime" + "strings" + + tea "charm.land/bubbletea/v2" +) + +// This file implements the platform-specific clipboard-image read behind Ctrl+V +// (and Cmd+V) image paste, mirroring Claude Code: when the system clipboard holds +// a raster image the model saves it to a temp PNG and drops an "[Image #N]" +// placeholder into the composer (see model.handleImagePaste), expanded at submit +// into an "@image:" reference that BuildUserContent attaches as multimodal +// content. Text clipboards fall through to the normal OSC52 read. + +// clipboardImageMsg is the reply to a clipboard-image read attempt. When ok is +// true, path is the temp PNG the decoded image was written to; when ok is false +// the clipboard held no image (or no reader tool was available) and the caller +// falls back to a plain text read. +type clipboardImageMsg struct { + path string + ok bool +} + +// readClipboardImage is a tea.Cmd that tries to pull a raster image out of the +// system clipboard and save it as a PNG under the OS temp dir. It shells out to +// the platform's clipboard tool (macOS: osascript; Linux: wl-paste or xclip). Any +// failure — no image on the clipboard, a missing tool — yields ok=false so the +// caller can fall back to a normal text paste rather than surfacing an error. +func readClipboardImage() tea.Msg { + path, ok := saveClipboardImage() + return clipboardImageMsg{path: path, ok: ok} +} + +// saveClipboardImage writes the clipboard image to a fresh temp PNG and returns +// its path, or ok=false when the clipboard holds no image. The temp file is +// removed on any failure so a stray empty file is never left behind. +func saveClipboardImage() (string, bool) { + f, err := os.CreateTemp("", "pigo-clip-*.png") + if err != nil { + return "", false + } + path := f.Name() + f.Close() + + var okRead bool + switch runtime.GOOS { + case "darwin": + okRead = saveClipboardImageDarwin(path) + case "linux": + okRead = saveClipboardImageLinux(path) + } + if !okRead { + os.Remove(path) + return "", false + } + if fi, err := os.Stat(path); err != nil || fi.Size() == 0 { + os.Remove(path) + return "", false + } + return path, true +} + +// darwinClipboardScript asks the pasteboard for its contents as PNG data and +// writes the raw bytes to the path passed as the first argv item, returning "ok" +// on success or "noimage" when the clipboard cannot be coerced to an image. +const darwinClipboardScript = `on run argv + set outPath to item 1 of argv + try + set pngData to (the clipboard as «class PNGf») + on error + return "noimage" + end try + set fh to open for access (POSIX file outPath) with write permission + set eof fh to 0 + write pngData to fh + close access fh + return "ok" +end run` + +// saveClipboardImageDarwin uses osascript (always present on macOS) to coerce the +// pasteboard to PNG and write it to path. +func saveClipboardImageDarwin(path string) bool { + out, err := exec.Command("osascript", "-e", darwinClipboardScript, path).Output() + return err == nil && strings.TrimSpace(string(out)) == "ok" +} + +// saveClipboardImageLinux reads an image/png off the clipboard via wl-paste +// (Wayland) or xclip (X11), preferring whichever is installed. It first checks the +// advertised MIME types so a text-only clipboard is not mistaken for an image. +func saveClipboardImageLinux(path string) bool { + if _, err := exec.LookPath("wl-paste"); err == nil { + types, _ := exec.Command("wl-paste", "--list-types").Output() + if strings.Contains(string(types), "image/png") { + if writeCmdOutput(path, exec.Command("wl-paste", "--type", "image/png")) { + return true + } + } + } + if _, err := exec.LookPath("xclip"); err == nil { + targets, _ := exec.Command("xclip", "-selection", "clipboard", "-t", "TARGETS", "-o").Output() + if strings.Contains(string(targets), "image/png") { + if writeCmdOutput(path, exec.Command("xclip", "-selection", "clipboard", "-t", "image/png", "-o")) { + return true + } + } + } + return false +} + +// writeCmdOutput runs cmd and writes its stdout to path, reporting whether it +// produced any bytes. +func writeCmdOutput(path string, cmd *exec.Cmd) bool { + out, err := cmd.Output() + if err != nil || len(out) == 0 { + return false + } + return os.WriteFile(path, out, 0o600) == nil +} diff --git a/pigo/internal/cli/tui/doc.go b/pigo/internal/cli/tui/doc.go new file mode 100644 index 0000000..7aaae0a --- /dev/null +++ b/pigo/internal/cli/tui/doc.go @@ -0,0 +1,14 @@ +// Package tui hosts the full-screen terminal UI for pigo's interactive mode +// (US-001). It is the alt-screen counterpart to the line-based REPL in +// internal/cli/repl: cmd/pigo's dispatch launches it via Run when there is no +// prompt, stdout is a TTY, and --no-tui is not set; otherwise the REPL path is +// used. See tasks/spec-tui-agent.md (Sections 2.1, 4.2, 5.2) for the design. +// +// This node is the skeleton: a root Model (Init/Update/View) built on Bubble +// Tea v2 (charm.land/bubbletea/v2) that renders an empty shell — a placeholder +// status bar, an empty transcript area, and an empty input line — starts on the +// alt-screen, and quits cleanly on Ctrl+C / Ctrl+D, restoring the terminal. +// Session assembly, the run bridge, tool cards, and slash-command completion +// land in downstream nodes; Options already mirrors repl.Options so those nodes +// can wire real behavior without changing the entry seam. +package tui diff --git a/pigo/internal/cli/tui/gitinfo.go b/pigo/internal/cli/tui/gitinfo.go new file mode 100644 index 0000000..8d9a43c --- /dev/null +++ b/pigo/internal/cli/tui/gitinfo.go @@ -0,0 +1,136 @@ +package tui + +import ( + "os/exec" + "strings" + + tea "charm.land/bubbletea/v2" +) + +// gitInfoMsg is the result of an async git probe (fetchGitCmd). It is defined +// here rather than in msgs.go to avoid conflicting with the event-bridge message +// set (#387). The status bar consumes it to render the branch + working-tree +// state segment. +// +// - branch: the current branch name (empty when detached / unknown). +// - ahead: commits the branch is ahead of its upstream (0 when unknown or no +// upstream). Derived cheaply from `git status --porcelain -b`. +// - dirty: number of changed/untracked entries reported by `git status +// --porcelain` (staged, unstaged, and untracked all count). +// - ok: false when the cwd is not a git repository or any git command +// failed; the status bar hides the git segment in that case. +type gitInfoMsg struct { + branch string + ahead int + dirty int + ok bool +} + +// fetchGitCmd returns a tea.Cmd that probes the git working tree rooted at cwd +// and reports a gitInfoMsg. It runs read-only git plumbing with fixed arguments +// (no user interpolation, so no command-injection surface) off the tea +// goroutine. Any error — not a repo, git missing, detached parse failure — +// collapses to gitInfoMsg{ok:false}, which the status bar renders as "no git". +func fetchGitCmd(cwd string) tea.Cmd { + return func() tea.Msg { + // One porcelain call with the branch header gives us branch name, ahead + // count, and every dirty/untracked entry in a single stable, parseable + // format (-z would drop the header line, so we keep the newline form). + out, err := runGit(cwd, "status", "--porcelain", "-b") + if err != nil { + return gitInfoMsg{ok: false} + } + return parseGitStatus(out) + } +} + +// runGit executes a read-only git subcommand in dir and returns its stdout. The +// argument list is always caller-controlled constants (see fetchGitCmd), never +// user input. +func runGit(dir string, args ...string) (string, error) { + cmd := exec.Command("git", args...) + cmd.Dir = dir + out, err := cmd.Output() + if err != nil { + return "", err + } + return string(out), nil +} + +// parseGitStatus parses the output of `git status --porcelain -b` into a +// gitInfoMsg. The first line is the branch header ("## branch...upstream [ahead N, +// behind M]"); every subsequent non-empty line is one changed or untracked entry. +// A parsed result always has ok=true, since reaching this point means git ran. +func parseGitStatus(out string) gitInfoMsg { + info := gitInfoMsg{ok: true} + lines := strings.Split(out, "\n") + for i, line := range lines { + if i == 0 { + info.branch, info.ahead = parseBranchHeader(line) + continue + } + if strings.TrimSpace(line) == "" { + continue + } + info.dirty++ + } + return info +} + +// parseBranchHeader extracts the branch name and ahead count from a porcelain +// branch header line, e.g.: +// +// ## master...origin/master [ahead 4, behind 1] +// ## feature-x +// ## HEAD (no branch) +// +// It returns the branch name and ahead count (0 when absent). A line that is not +// a branch header yields ("", 0). +func parseBranchHeader(line string) (string, int) { + const prefix = "## " + if !strings.HasPrefix(line, prefix) { + return "", 0 + } + rest := strings.TrimPrefix(line, prefix) + + // Split off the optional " [ahead N, behind M]" tracking suffix. + branchPart := rest + ahead := 0 + if idx := strings.Index(rest, " ["); idx >= 0 { + branchPart = rest[:idx] + ahead = parseAhead(rest[idx:]) + } + + // "## HEAD (no branch)" — detached; keep the raw token as the branch label. + branchPart = strings.TrimSpace(branchPart) + + // Trim the "...upstream" tracking-branch tail if present. + if idx := strings.Index(branchPart, "..."); idx >= 0 { + branchPart = branchPart[:idx] + } + return branchPart, ahead +} + +// parseAhead pulls the integer following "ahead " out of a tracking suffix such +// as "[ahead 4, behind 1]". It returns 0 when no ahead count is present. +func parseAhead(suffix string) int { + const marker = "ahead " + idx := strings.Index(suffix, marker) + if idx < 0 { + return 0 + } + digits := suffix[idx+len(marker):] + n := 0 + found := false + for _, r := range digits { + if r < '0' || r > '9' { + break + } + n = n*10 + int(r-'0') + found = true + } + if !found { + return 0 + } + return n +} diff --git a/pigo/internal/cli/tui/gitinfo_test.go b/pigo/internal/cli/tui/gitinfo_test.go new file mode 100644 index 0000000..e9007e9 --- /dev/null +++ b/pigo/internal/cli/tui/gitinfo_test.go @@ -0,0 +1,99 @@ +package tui + +import "testing" + +func TestParseGitStatusDirtyCount(t *testing.T) { + // Sample `git status --porcelain -b` output: header + 4 entries (staged, + // modified, deleted, untracked). + out := "## master...origin/master [ahead 4, behind 1]\n" + + "M cmd/pigo/main.go\n" + + " M internal/foo.go\n" + + "D cmd/pigo/run.go\n" + + "?? new_file.go\n" + + info := parseGitStatus(out) + if !info.ok { + t.Fatal("parseGitStatus should report ok=true") + } + if info.branch != "master" { + t.Errorf("branch = %q, want master", info.branch) + } + if info.dirty != 4 { + t.Errorf("dirty = %d, want 4", info.dirty) + } + if info.ahead != 4 { + t.Errorf("ahead = %d, want 4", info.ahead) + } +} + +func TestParseGitStatusCleanTree(t *testing.T) { + out := "## main...origin/main\n" + info := parseGitStatus(out) + if !info.ok { + t.Fatal("ok should be true") + } + if info.branch != "main" { + t.Errorf("branch = %q, want main", info.branch) + } + if info.dirty != 0 { + t.Errorf("dirty = %d, want 0", info.dirty) + } + if info.ahead != 0 { + t.Errorf("ahead = %d, want 0", info.ahead) + } +} + +func TestParseGitStatusNoUpstream(t *testing.T) { + out := "## feature-x\n M a.go\n" + info := parseGitStatus(out) + if info.branch != "feature-x" { + t.Errorf("branch = %q, want feature-x", info.branch) + } + if info.ahead != 0 { + t.Errorf("ahead = %d, want 0", info.ahead) + } + if info.dirty != 1 { + t.Errorf("dirty = %d, want 1", info.dirty) + } +} + +func TestParseGitStatusDetached(t *testing.T) { + out := "## HEAD (no branch)\n" + info := parseGitStatus(out) + if info.branch != "HEAD (no branch)" { + t.Errorf("branch = %q, want %q", info.branch, "HEAD (no branch)") + } +} + +func TestParseBranchHeaderAhead(t *testing.T) { + branch, ahead := parseBranchHeader("## dev...origin/dev [ahead 12]") + if branch != "dev" { + t.Errorf("branch = %q, want dev", branch) + } + if ahead != 12 { + t.Errorf("ahead = %d, want 12", ahead) + } +} + +func TestParseBranchHeaderBehindOnly(t *testing.T) { + branch, ahead := parseBranchHeader("## dev...origin/dev [behind 3]") + if branch != "dev" { + t.Errorf("branch = %q, want dev", branch) + } + if ahead != 0 { + t.Errorf("ahead = %d, want 0 (behind only)", ahead) + } +} + +func TestFetchGitCmdNonRepo(t *testing.T) { + // A path that is not a git repository should collapse to ok=false. /tmp is + // (almost) never a repo; if it somehow is on this host, skip. + msg := fetchGitCmd("/")() + git, ok := msg.(gitInfoMsg) + if !ok { + t.Fatalf("expected gitInfoMsg, got %T", msg) + } + if git.ok { + t.Skip("host root unexpectedly reports a git repo; skipping") + } +} diff --git a/pigo/internal/cli/tui/host.go b/pigo/internal/cli/tui/host.go new file mode 100644 index 0000000..f9b5cbd --- /dev/null +++ b/pigo/internal/cli/tui/host.go @@ -0,0 +1,91 @@ +// This file makes runSession satisfy cli.Host: the accessor and mutator methods +// let the /status command (and future /goal, /btw) read the session's live +// collaborators and mutable state through the cli.Host contract rather than the +// concrete aggregate — the same seam the REPL's replDeps implements. The +// compile-time assertion below fails the build if runSession drifts out of +// conformance. +package tui + +import ( + "bufio" + "fmt" + "io" + "sync" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +var _ cli.Host = (*runSession)(nil) + +func (s *runSession) Store() *session.Store { return s.store } +func (s *runSession) Header() session.SessionHeader { return s.header } +func (s *runSession) AgentCtx() *agentcore.AgentContext { return s.agentCtx } +func (s *runSession) Live() *cli.LiveConfig { return s.live } +func (s *runSession) Registry() *agenttool.ToolRegistry { return s.reg } +func (s *runSession) Reminders() *runtime.ReminderRegistry { return s.reminders } +func (s *runSession) Slash() *runtime.SlashRegistry { return s.slash } +func (s *runSession) Creds() *provider.CredentialStore { return s.creds } +func (s *runSession) Notifier() *plugin.EventNotifier { return nil } +func (s *runSession) NotifierHandle() func(agentcore.AgentEvent) { return s.onEvent } +func (s *runSession) Trust() *trust.Manager { return s.trust } +func (s *runSession) Goal() *agenttool.GoalState { return nil } +func (s *runSession) Telemetry() *cli.TelemetryHolder { return s.telemetry } +func (s *runSession) Dispatcher() *hooks.Dispatcher { return s.dispatcher } +func (s *runSession) HookDeps() run.HookDeps { return s.hookDeps } +func (s *runSession) Cwd() string { return s.cwd } +func (s *runSession) Input() *bufio.Reader { return nil } +func (s *runSession) ConfirmMu() *sync.Mutex { return nil } + +func (s *runSession) CurLeaf() string { return s.curLeaf } +func (s *runSession) SetCurLeaf(id string) { s.curLeaf = id } +func (s *runSession) Persisted() int { return s.persisted } +func (s *runSession) SetPersisted(n int) { s.persisted = n } + +func (s *runSession) LastBtw() *agentcore.AgentContext { return s.lastBtw } +func (s *runSession) SetLastBtw(ctx *agentcore.AgentContext) { s.lastBtw = ctx } +func (s *runSession) LastBtwBase() int { return s.lastBtwBase } +func (s *runSession) SetLastBtwBase(n int) { s.lastBtwBase = n } + +// renderSession writes the /session summary (US-009, #125) to out — the same +// format the REPL's runSession prints: session id, message count, estimated +// token usage, model/provider, creation time, and compaction-checkpoint count. +// It lives on runSession so the TUI's /session intercept and the REPL share one +// rendering; counts derive from the in-memory context (the source of truth for +// the live turn), so unsaved messages are counted too. +func (s *runSession) renderSession(out io.Writer) { + msgs := s.agentCtx.Messages + tokens := compaction.EstimateContextTokens(msgs).Tokens + compactions := 0 + for _, m := range msgs { + if _, ok := m.(agentcore.CompactionMessage); ok { + compactions++ + } + } + fmt.Fprintf(out, "session: %s\n", s.header.ID) + fmt.Fprintf(out, "messages: %d\n", len(msgs)) + fmt.Fprintf(out, "tokens (est): %d\n", tokens) + model := s.live.Model + providerName := s.live.ProviderName + if model == "" { + model = s.header.Model + } + if providerName == "" { + providerName = s.header.Provider + } + fmt.Fprintf(out, "model: %s (provider: %s)\n", model, providerName) + if !s.header.CreatedAt.IsZero() { + fmt.Fprintf(out, "created: %s\n", s.header.CreatedAt.Format(time.RFC3339)) + } + fmt.Fprintf(out, "compactions: %d\n", compactions) +} diff --git a/pigo/internal/cli/tui/input.go b/pigo/internal/cli/tui/input.go new file mode 100644 index 0000000..2c91af6 --- /dev/null +++ b/pigo/internal/cli/tui/input.go @@ -0,0 +1,167 @@ +package tui + +import ( + "charm.land/bubbles/v2/key" + "charm.land/bubbles/v2/textarea" + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" +) + +// This file implements the prompt input field of the full-screen TUI (US-007, +// FR-11/13/14). It wraps charm.land/bubbles/v2/textarea into a small `input` +// component so the model can embed a real multi-line editor instead of the +// throwaway string buffer the skeleton shipped with. +// +// Why textarea rather than a hand-rolled buffer: textarea edits by grapheme / +// rune, so CJK and emoji are inserted and deleted whole. This is exactly the +// class of bug the old REPL input had — it keyed on byte length (len==1) and +// silently dropped the trailing bytes of every multi-byte rune. We deliberately +// delegate all character handling to textarea and never touch bytes ourselves. +// +// Shift+Enter inserts a newline so the editor is a true multi-line composer; +// plain Enter submits (intercepted by the model, never reaching textarea). The +// default InsertNewline binding (Enter) is therefore rebound to Shift+Enter. See +// model.handleKey. + +// maxInputRows caps how tall the editor grows as the user adds lines. Past this +// the buffer keeps growing but textarea scrolls its own viewport, so the shell's +// row accounting stays bounded and the transcript never collapses to nothing. +const maxInputRows = 6 + +// input is the prompt editor. It embeds a textarea.Model and exposes just the +// surface the root model needs: value/clear, focus/blur (input is blurred while +// a run is in flight so keystrokes never corrupt an in-flight prompt), a width +// setter driven by tea.WindowSizeMsg, and a render string for View. +type input struct { + ta textarea.Model + // width is the full editor width (terminal columns) last set via SetWidth. It + // is the span of the top/bottom rules drawn around the editor (Claude-Code + // style), kept separately because textarea's own Width() reports only the + // inner text area (prompt column excluded). + width int +} + +// newInput builds a focused editor. It starts one row tall and grows with the +// buffer (up to maxInputRows) as the user inserts newlines with Enter. The +// buffer itself is unbounded; beyond maxInputRows textarea scrolls internally. +func newInput() input { + ta := textarea.New() + ta.Prompt = "> " + ta.Placeholder = "Type a message… (Enter to send, Shift+Enter for newline)" + ta.ShowLineNumbers = false + ta.CharLimit = 0 + // Let the textarea own its own height: DynamicHeight grows/shrinks it to the + // content between MinHeight (1) and MaxHeight (maxInputRows), and — critically + // — fixes the viewport scroll offset in the same pass. Doing it manually (an + // after-the-fact SetHeight in syncHeight) left a stale scroll offset: inserting + // a newline scrolled the cursor into view while the editor was still 1 row + // tall, pushing the first line off the top, and the later SetHeight never + // scrolled it back — so a two-line buffer rendered as two blank lines. + ta.MinHeight = 1 + ta.MaxHeight = maxInputRows + ta.DynamicHeight = true + // textarea.New starts at defaultHeight (6). DynamicHeight only recomputes on + // edits, so pin the empty editor to one row up front — otherwise the shell + // would reserve six rows before the user has typed anything. + ta.SetHeight(1) + // Rebind InsertNewline from its default (Enter) to the newline keys, since + // plain Enter is the model's submit key (handleKey intercepts it before + // textarea sees it). Shift+Enter is the primary, advertised binding: Bubble + // Tea v2 already enables the Kitty keyboard protocol's disambiguate flag + // (flag 1) on every View, so capable terminals (kitty, ghostty, wezterm, + // recent iTerm2) report Shift+Enter as a distinct CSI-u sequence rather than + // a bare CR. Crucially this is flag 1, NOT flag 8 (ReportAllKeysAsEscapeCodes) + // — flag 8 broke IME / CJK input because it strips associated text, whereas + // flag 1 only disambiguates special keys and leaves text entry untouched. + // On terminals without the protocol (macOS Terminal.app, tmux by default) + // Shift+Enter arrives byte-identical to Enter and would submit, so Ctrl+J (a + // literal LF, always distinct from Enter's CR) and Alt+Enter (ESC-prefixed, + // always distinct) are kept as silent fallbacks — a newline is guaranteed to + // work everywhere. All three split the line at the cursor and keep typed text. + ta.KeyMap.InsertNewline = key.NewBinding( + key.WithKeys("shift+enter", "ctrl+j", "alt+enter"), + key.WithHelp("shift+enter", "insert newline"), + ) + // Draw the cursor into the rendered string: the model composes View as a + // plain string rather than driving textarea's real cursor reporting. + ta.SetVirtualCursor(true) + // Drop the default cursor-line background highlight so the composer is framed + // only by the top/bottom rules (see View), matching Claude Code — no fill. + styles := ta.Styles() + styles.Focused.CursorLine = lipgloss.NewStyle() + styles.Blurred.CursorLine = lipgloss.NewStyle() + ta.SetStyles(styles) + ta.Focus() + return input{ta: ta} +} + +// Update forwards a message (typically a key press) to the underlying textarea +// and returns the updated component. The model calls this only for keys it does +// not intercept itself (submit / interrupt / quit), so textarea sees ordinary +// editing keys — including Enter (newline) and CJK / emoji runes, which it +// inserts whole. Height is owned by textarea's DynamicHeight (see newInput), so +// there is nothing to re-sync here. +func (in input) Update(msg tea.Msg) (input, tea.Cmd) { + var cmd tea.Cmd + in.ta, cmd = in.ta.Update(msg) + return in, cmd +} + +// Height reports the current visible row count of the editor so the model can +// reserve that many rows in its View layout. It includes the two rule rows (top +// and bottom) drawn around the textarea. +func (in input) Height() int { return in.ta.Height() + 2 } + +// Value returns the current buffer contents, including any embedded newlines. +func (in input) Value() string { return in.ta.Value() } + +// SetValue replaces the buffer contents and moves the cursor to the end. It is +// used by slash autocomplete (Tab) to complete the buffer to the chosen command. +func (in *input) SetValue(s string) { + in.ta.SetValue(s) +} + +// Clear empties the buffer and resets the cursor to the start. +func (in *input) Clear() { + in.ta.Reset() +} + +// Focus enables editing and returns the cursor-blink Cmd. +func (in *input) Focus() tea.Cmd { return in.ta.Focus() } + +// Blur disables editing (used while a run is in flight). +func (in *input) Blur() { in.ta.Blur() } + +// Focused reports whether the editor currently accepts input. +func (in input) Focused() bool { return in.ta.Focused() } + +// Line reports the zero-based index of the line the cursor is on, and LineCount +// the total number of lines in the buffer. The model uses them to decide whether +// ↑/↓ should walk the prompt history (caret on the first / last line) or move the +// caret within a multi-line draft. +func (in input) Line() int { return in.ta.Line() } +func (in input) LineCount() int { return in.ta.LineCount() } + +// SetWidth resizes the editor to the terminal width so wrapping and the prompt +// column line up with the rest of the shell. +func (in *input) SetWidth(w int) { + if w < 0 { + w = 0 + } + in.width = w + in.ta.SetWidth(w) +} + +// View renders the editor to a string for embedding in the model's View. The +// textarea is framed with a top and bottom rule (no side borders) in the muted +// gray, mirroring Claude Code's composer — a pair of horizontal lines rather +// than a background fill. The rules span the full editor width. +func (in input) View() string { + style := lipgloss.NewStyle(). + Border(lipgloss.NormalBorder(), true, false, true, false). + BorderForeground(lipgloss.Color(colorGray)) + if in.width > 0 { + style = style.Width(in.width) + } + return style.Render(in.ta.View()) +} diff --git a/pigo/internal/cli/tui/input_test.go b/pigo/internal/cli/tui/input_test.go new file mode 100644 index 0000000..4fe592d --- /dev/null +++ b/pigo/internal/cli/tui/input_test.go @@ -0,0 +1,205 @@ +package tui + +import ( + "strings" + "testing" + "unicode/utf8" + + tea "charm.land/bubbletea/v2" +) + +// runeKey builds a printable-character key press carrying r, mirroring what a +// terminal sends for a typed rune: Code is the rune and Text is its UTF-8 +// encoding (textarea inserts from Text). This is how CJK / emoji reach the +// component. +func runeKey(r rune) tea.KeyPressMsg { + return tea.KeyPressMsg{Code: r, Text: string(r)} +} + +// TestInputCJKByRune drives the input with the runes of "你好" and asserts the +// buffer holds the full multi-byte string with the cursor left on a rune +// boundary. This guards against the old REPL bug that keyed on byte length +// (len==1) and dropped the trailing bytes of every multi-byte rune. +func TestInputCJKByRune(t *testing.T) { + in := newInput() + for _, r := range "你好" { + var cmd tea.Cmd + in, cmd = in.Update(runeKey(r)) + _ = cmd + } + + got := in.Value() + if got != "你好" { + t.Fatalf("Value() = %q, want %q", got, "你好") + } + if !utf8.ValidString(got) { + t.Fatalf("Value() is not valid UTF-8: %q", got) + } + if n := utf8.RuneCountInString(got); n != 2 { + t.Fatalf("rune count = %d, want 2 (no dropped chars)", n) + } + // The cursor column is a rune index into the line; after two runes it must be + // 2, proving textarea advanced by whole runes rather than bytes. + if col := in.ta.Column(); col != 2 { + t.Errorf("cursor column = %d, want 2 (rune boundary)", col) + } +} + +// TestInputEmojiByRune confirms a multi-byte emoji is inserted whole. +func TestInputEmojiByRune(t *testing.T) { + in := newInput() + in, _ = in.Update(runeKey('🚀')) + if got := in.Value(); got != "🚀" { + t.Fatalf("Value() = %q, want %q", got, "🚀") + } +} + +// TestInputNewlineKeys verifies each newline binding — Shift+Enter (primary), +// Ctrl+J and Alt+Enter (fallbacks) — inserts a newline into the buffer while +// plain runes fill each line. +func TestInputNewlineKeys(t *testing.T) { + cases := []struct { + name string + key tea.KeyPressMsg + }{ + {"shift+enter", tea.KeyPressMsg{Code: tea.KeyEnter, Mod: tea.ModShift}}, + {"ctrl+j", tea.KeyPressMsg{Code: 'j', Mod: tea.ModCtrl}}, + {"alt+enter", tea.KeyPressMsg{Code: tea.KeyEnter, Mod: tea.ModAlt}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + in := newInput() + in, _ = in.Update(runeKey('你')) + in, _ = in.Update(tc.key) + in, _ = in.Update(runeKey('好')) + if got := in.Value(); got != "你\n好" { + t.Fatalf("Value() = %q, want %q", got, "你\n好") + } + }) + } +} + +// TestInputEnterIsNotNewline confirms plain Enter does NOT insert a newline in +// the editor: the model intercepts it as submit, so the editor must leave it +// alone (only Shift+Enter breaks a line). +func TestInputEnterIsNotNewline(t *testing.T) { + in := newInput() + in, _ = in.Update(runeKey('a')) + in, _ = in.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if got := in.Value(); got != "a" { + t.Fatalf("Value() = %q, want %q (Enter must not add a newline)", got, "a") + } +} + +// TestInputNewlineRendersBothLines guards the viewport-offset regression: after +// 你 + Shift+Enter + 好 the buffer is "你\n好", but the editor once rendered two +// blank rows because a manual SetHeight left the textarea's viewport scrolled +// past the first line. DynamicHeight now resets the scroll offset in the same +// pass, so both runes must appear in the rendered View. +func TestInputNewlineRendersBothLines(t *testing.T) { + in := newInput() + in.SetWidth(40) + in, _ = in.Update(runeKey('你')) + in, _ = in.Update(tea.KeyPressMsg{Code: tea.KeyEnter, Mod: tea.ModShift}) + in, _ = in.Update(runeKey('好')) + + view := in.View() + if !strings.Contains(view, "你") || !strings.Contains(view, "好") { + t.Fatalf("rendered view missing content, want both 你 and 好:\n%s", view) + } +} + +// TestInputClearBlurFocus exercises the lifecycle methods the model relies on +// while gating input during a run. +func TestInputClearBlurFocus(t *testing.T) { + in := newInput() + in, _ = in.Update(runeKey('x')) + in.Clear() + if got := in.Value(); got != "" { + t.Errorf("after Clear, Value() = %q, want empty", got) + } + if !in.Focused() { + t.Error("newInput should start focused") + } + in.Blur() + if in.Focused() { + t.Error("after Blur, Focused() should be false") + } + in.Focus() + if !in.Focused() { + t.Error("after Focus, Focused() should be true") + } +} + +// TestModelEnterSubmits feeds a typed line and Enter to the model and asserts +// the prompt is submitted: the user turn lands in the transcript and the editor +// is cleared. Enter submits; Shift+Enter is the newline key in the multi-line +// composer. With no startRunFn wired the model stays idle and records the +// pre-#392 system note. +func TestModelEnterSubmits(t *testing.T) { + m := NewModel(Options{}) + var model tea.Model = m + for _, r := range "你好世界" { + model, _ = model.Update(runeKey(r)) + } + model, _ = model.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + + got := model.(Model) + if got.input.Value() != "" { + t.Errorf("after submit, input = %q, want cleared", got.input.Value()) + } + joined := strings.Join(blockTexts(got.transcript), "\n") + if !strings.Contains(joined, "你好世界") { + t.Errorf("submitted prompt missing from transcript: %q", joined) + } +} + +// TestModelTwoStageInterrupt verifies FR-14: while running, Esc / Ctrl+C +// interrupts the in-flight run (calls interruptFn) and does NOT quit; while +// idle, the same keys quit the program. +func TestModelTwoStageInterrupt(t *testing.T) { + for _, key := range []tea.KeyPressMsg{ + {Code: tea.KeyEscape}, + {Code: 'c', Mod: tea.ModCtrl}, + } { + // Running: first press interrupts, no quit. + interrupted := false + running := NewModel(Options{}) + running.running = true + running.interruptFn = func() { interrupted = true } + next, cmd := running.Update(key) + if !interrupted { + t.Errorf("%s while running: interruptFn was not called", key.String()) + } + if next.(Model).quitting { + t.Errorf("%s while running: model should not be quitting", key.String()) + } + if cmd != nil { + if _, isQuit := cmd().(tea.QuitMsg); isQuit { + t.Errorf("%s while running: should not quit", key.String()) + } + } + + // Idle: the same key quits. + idle := NewModel(Options{}) + got, cmd := idle.Update(key) + if cmd == nil { + t.Fatalf("%s while idle: expected a quit command", key.String()) + } + if _, isQuit := cmd().(tea.QuitMsg); !isQuit { + t.Errorf("%s while idle: cmd should be tea.Quit", key.String()) + } + if !got.(Model).quitting { + t.Errorf("%s while idle: model should be marked quitting", key.String()) + } + } +} + +// blockTexts extracts the raw text of every transcript block for assertions. +func blockTexts(t transcript) []string { + out := make([]string, len(t.blocks)) + for i, b := range t.blocks { + out[i] = b.text + } + return out +} diff --git a/pigo/internal/cli/tui/markdown.go b/pigo/internal/cli/tui/markdown.go new file mode 100644 index 0000000..dbd4a3f --- /dev/null +++ b/pigo/internal/cli/tui/markdown.go @@ -0,0 +1,122 @@ +package tui + +import ( + "strings" + "sync" + + "github.com/charmbracelet/glamour" + "github.com/charmbracelet/x/ansi" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// This file renders finalized assistant turns as Markdown inside the TUI +// transcript (fix #3, mirroring the REPL's ui.RenderMarkdown). The REPL renders +// once at turn-end because Markdown can only be laid out when the whole block is +// known; the transcript does the same — only a finalized assistant block is +// passed through here, never the still-streaming one. +// +// Unlike the REPL's shared renderer (WithWordWrap(0), which relies on the raw +// terminal to soft-wrap), the transcript lives inside a fixed-width viewport +// that does NOT soft-wrap, so we must wrap the Markdown to the content width +// ourselves. Renderers are therefore cached per width and rebuilt when the width +// changes (a resize), which is rare enough that the rebuild cost is negligible. + +var ( + mdMu sync.Mutex + mdCache = map[int]*glamour.TermRenderer{} + // mdDark selects the glamour style: a dark palette when true (the default, + // matching most terminals), a light palette when false. It is set once from + // the terminal's real background via SetMarkdownDark and never queried at + // render time — see the comment there. + mdDark = true +) + +// SetMarkdownDark records whether the terminal has a dark background and drops +// the renderer cache so the next render rebuilds with the matching style. +// +// This is the fix for the escape-sequence leak into the input box: glamour's +// WithAutoStyle() detects the palette by issuing its OWN synchronous OSC 11 +// background-color query and reading the reply straight from the tty. Under the +// alt-screen, bubbletea already owns the input reader, so that reply +// (\x1b]11;rgb:…\x07) races with — and is swallowed by — bubbletea's parser, +// which then leaks the unparsed tail (e.g. "1;rgb:0000/0000/0000" plus a stray +// SGR mouse report) into the textarea as literal text. We instead let bubbletea +// detect the background the parser-safe way (RequestBackgroundColor → +// BackgroundColorMsg) and feed the result here, then build glamour with a fixed +// WithStandardStyle so it never touches the terminal. +func SetMarkdownDark(dark bool) { + mdMu.Lock() + defer mdMu.Unlock() + if dark == mdDark { + return + } + mdDark = dark + mdCache = map[int]*glamour.TermRenderer{} +} + +// rendererFor returns a glamour renderer that word-wraps to width columns, +// building and caching one per distinct width. A build failure caches nothing +// and returns nil so callers fall back to the raw source. The style is fixed +// (WithStandardStyle) rather than auto-detected, so building a renderer never +// queries the terminal. +func rendererFor(width int) *glamour.TermRenderer { + mdMu.Lock() + defer mdMu.Unlock() + if r, ok := mdCache[width]; ok { + return r + } + wrap := width + if wrap < 0 { + wrap = 0 + } + style := "dark" + if !mdDark { + style = "light" + } + r, err := glamour.NewTermRenderer( + glamour.WithStandardStyle(style), + glamour.WithWordWrap(wrap), + ) + if err != nil { + return nil + } + mdCache[width] = r + return r +} + +// renderMarkdown returns src rendered as styled terminal Markdown wrapped to +// width columns. It is gated exactly like the REPL's renderer: when output is +// not an interactive terminal (pipes, tests) the raw source is returned so +// golden tests and machine consumers are unaffected. A nil/broken renderer or a +// render error also returns the raw source, so content is never dropped. The +// trailing newline glamour appends is trimmed so the block joins cleanly with +// its neighbors in the transcript. +func renderMarkdown(src string, width int) string { + if !ui.Enabled() { + return src + } + if strings.TrimSpace(src) == "" { + return src + } + r := rendererFor(width) + if r == nil { + return src + } + out, err := r.Render(src) + if err != nil { + return src + } + // Glamour word-wraps prose to width, but its document margin and + // non-wrapping elements (code blocks, tables) can still emit lines wider than + // the content column. The transcript viewport does not clip horizontally, so + // an over-wide line would spill into (and visually erase) the persistent + // scrollbar column on its right. Hard-wrap the rendered output — ANSI- and + // wide-char-aware — so every line fits within width and the scrollbar stays + // put. Prose already within width is untouched. + trimmed := strings.Trim(out, "\n") + if width > 0 { + trimmed = ansi.Hardwrap(trimmed, width, false) + } + return trimmed +} diff --git a/pigo/internal/cli/tui/model.go b/pigo/internal/cli/tui/model.go new file mode 100644 index 0000000..8a671b4 --- /dev/null +++ b/pigo/internal/cli/tui/model.go @@ -0,0 +1,1360 @@ +package tui + +import ( + "bytes" + "fmt" + "os" + "regexp" + "strconv" + "strings" + "time" + + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/memstatus" + "github.com/smallnest/pigo/internal/cli/status" + "github.com/smallnest/pigo/internal/memory" + "github.com/smallnest/pigo/internal/runtime" +) + +// Model is the root Bubble Tea model for the full-screen TUI. It composes a +// scrolling transcript (US-005) with the persistent status bar (#386) and a +// minimal input line, and owns the run lifecycle: on prompt submit it starts an +// agent run through the event bridge (bridge.go) and pumps the resulting tea.Msg +// stream into the transcript one message at a time. Downstream nodes grow the +// input into a full textarea (#390), render tool cards (#389), and wire the real +// session/run assembly (#392); the Init/Update/View contract and the alt-screen +// + quit-key handling stay stable. +type Model struct { + opts Options + theme Theme + + // width and height track the terminal size reported by tea.WindowSizeMsg. + // They are zero until the first size message arrives; View degrades to a + // minimal render in that window. + width int + height int + + // transcript is the scrolling message log (user / assistant / system turns). + transcript transcript + + // input is the multi-line prompt editor (#390). It wraps a bubbles textarea + // so CJK / emoji are edited by rune (no dropped-byte bug), Enter submits and + // Shift+Enter inserts a newline. It is blurred while a run is in flight. + input input + + // history holds previously submitted inputs (prompts and slash commands, in + // order), and histIdx is the browse cursor into it: len(history) means "not + // browsing — on the live draft", any smaller index points at a recalled entry. + // histDraft stashes the in-progress buffer when browsing begins so ↓ past the + // newest entry restores it. ↑/↓ walk history when the caret is on the first / + // last line of the composer, so multi-line editing is unaffected. + history []string + histIdx int + histDraft string + + // running is true while an agent run is draining through runCh. Input submit + // is gated on it so a new run cannot start mid-run. + running bool + // runCh is the bridge channel for the in-flight run, or nil when idle. Update + // re-issues waitForEvent(runCh) after every bridged msg except runEndMsg. + runCh chan tea.Msg + + // startRunFn launches an agent run for the submitted prompt, returning the + // bridge channel and the first waitForEvent Cmd (see bridge.startRun). It is + // bound to runSession.startRun by withSession (#392): the real binding + // constructs an AgentContext + RunConfig from opts and the live session. It is + // nil for a session-less model (the pure constructor / tests), in which case a + // submit records the prompt but starts no run. + startRunFn func(prompt string) (chan tea.Msg, tea.Cmd) + + // session is the assembled run/persistence state (store, header, growing + // context, live config). It is nil for a session-less model; when set, the + // model persists the conversation to ~/.pigo/sessions after each turn ends. + session *runSession + + // interruptFn cancels the in-flight run (the first stage of the two-stage + // interrupt, FR-14): pressing Esc / Ctrl+C while running signals the run to + // stop rather than quitting the program. It is a seam wired alongside + // startRunFn by session assembly (#392) — typically the run ctx's cancel + // func. Until then it may be nil, in which case an interrupt while running is + // a safe no-op (the pump keeps draining until it ends on its own). + interruptFn func() + + // quitting is set when a quit key (Ctrl+C / Ctrl+D) is seen, so View can be a + // no-op on the final frame while the program tears down and restores the + // terminal. + quitting bool + + // statusBar renders the persistent bottom line (#386, US-003). It is fed the + // terminal width, telemetry-derived context usage, and the async git probe + // result; View renders it just above the input line. + statusBar statusBar + + // cwd is the launch directory, captured once at construction and reused for + // the git probe and the status bar's path display. + cwd string + + // slash is the shared slash-command registry (#383) the TUI consults exactly + // as the REPL does: /model, /help, user templates, plugin commands and skills. + // It is bound to live so a /model switch mutates the same config the run loop + // reads. Built in NewModel (built-ins + disk templates) and rebuilt in + // withSession against the session's live config. + slash *runtime.SlashRegistry + // live is the mutable run configuration the /model command switches. In a + // session-bound model it is the SAME pointer the run loop reads (set by + // withSession), so a switch takes effect on the next turn. + live *cli.LiveConfig + // menu is the autocomplete popup shown while a "/name" is being typed (#391). + // It filters slash by the typed prefix; the model intercepts arrow/Tab/Enter + // keys to drive it before delegating to the textarea. + menu slashMenu + + // toolCards indexes the rich tool-call cards (#389, US-006) by tool-call id so + // a toolEndMsg can locate the card started earlier and flip its state / attach + // the parsed response. Each card is also appended to the transcript as an + // ordered block (by pointer), so mutating one here re-renders it inline on the + // next reflow. + toolCards map[string]*toolCard + // lastToolCard points at the most recently started card; Ctrl+O toggles its + // expanded state and re-flows the transcript. + lastToolCard *toolCard + + // draggingScrollbar is set while the left mouse button is held after pressing + // on the transcript scrollbar column, so subsequent motion events drag the + // thumb (and scroll the viewport) until the button is released. + draggingScrollbar bool + + // sel is the current mouse text selection over the rendered shell (screen + // cells). A left-press off the scrollbar starts it, drag extends it, and it + // persists after release so Ctrl+C can copy the highlighted text. + sel selection + + // spinner is the animated "working" indicator (verb + elapsed/token/effort + // stats) shown on the row above the input while a run is in flight. + spinner spinner + + // subagents is the ordered set of live sub-agents dispatched by the `task` + // tool (SPEC 4.4, US-006). A toolStartMsg with name=="task" adds a row (and + // records its start time), subagentProgressMsg refreshes activity/tokens, and + // the task's toolEndMsg removes it. View renders it as a multi-line panel just + // above the spinner; it contributes zero rows when empty. + subagents subagentPanel + + // pastes stores the full text of collapsed multi-line pastes, keyed by the id + // shown in the "[Pasted text #N +M lines]" placeholder left in the composer. + // submit expands the placeholders back to their content before sending, so a + // large paste never floods the editor (mirroring Claude Code). + pastes map[int]string + // pasteSeq is the monotonic counter behind the paste placeholder ids. It keeps + // climbing across submits so ids stay unique for the session. + pasteSeq int + + // images maps the id shown in an "[Image #N]" placeholder to the temp PNG a + // Ctrl+V / Cmd+V image paste was saved to. submit expands the placeholder to an + // "@image:" reference so BuildUserContent attaches the image as + // multimodal content (mirroring Claude Code's image paste). + images map[int]string + // imageSeq is the monotonic counter behind the image placeholder ids. + imageSeq int +} + +// NewModel builds the root model from the assembled Options. It reads the +// current working directory (for the status bar's path display and git probe) +// and assembles the shared slash-command registry (#391), which reads the user +// prompt-template dirs (~/.pigo/{commands,prompts}) and the pre-loaded skills; +// missing dirs are not an error. The registry is bound here to a live config +// derived from Options; withSession rebinds it to the session's live config so a +// /model switch reaches the run loop. +func NewModel(opts Options) Model { + cwd, err := os.Getwd() + if err != nil { + cwd = "" + } + theme := DefaultTheme() + live := &cli.LiveConfig{ + Model: opts.Model, + ProviderName: opts.ProviderName, + Provider: opts.Provider, + BaseURL: opts.BaseURL, + Protocol: opts.Protocol, + ThinkingLevel: opts.ThinkingLevel, + ContextWindow: cli.DefaultContextWindow, + } + return Model{ + opts: opts, + theme: theme, + transcript: newTranscript(theme), + input: newInput(), + cwd: cwd, + statusBar: newStatusBar(theme, opts, cwd), + toolCards: make(map[string]*toolCard), + slash: newSlashRegistry(opts, live), + live: live, + menu: newSlashMenu(theme), + spinner: newSpinner(theme), + pastes: make(map[int]string), + images: make(map[int]string), + } +} + +// withSession binds the assembled run session to the model: it wires the real +// run seam (startRunFn) and, for a resumed session, replays the prior history +// into the transcript so the user sees the conversation so far before entering +// interactive mode. Run calls it right after NewModel; the session-less +// constructor path (tests, pure construction) leaves startRunFn nil. +func (m Model) withSession(s *runSession, history []agentcore.Message) Model { + m.session = s + m.startRunFn = s.startRun + m.interruptFn = s.interrupt + // Rebind the registry to the session's own one (assembled against s.live, the + // very config the run loop reads via buildConfig) so /model mutates the live + // config, /trust reaches the session's trust manager (registered in + // newRunSessionWithStore), and /status can list skill/plugin/user commands. + m.live = s.live + m.slash = s.slash + m.transcript.addBanner(renderBanner(m.theme, m.opts, m.cwd)) + seedTranscript(&m.transcript, history) + return m +} + +// Init implements tea.Model. It kicks off the async git probe so the status bar +// can show the branch/dirty state as soon as it resolves; the alt-screen is +// requested declaratively via the AltScreen field on the View returned by View. +func (m Model) Init() tea.Cmd { + return tea.Batch(fetchGitCmd(m.cwd), m.input.Focus(), func() tea.Msg { + return tea.RequestBackgroundColor() + }) +} + +// Update implements tea.Model. It tracks the terminal size, drives the minimal +// input line, starts runs on submit, and pumps bridged run events into the +// transcript and status bar. It quits on the standard exit keys (Ctrl+C / +// Ctrl+D). +func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.height = msg.Height + m.relayout() + return m, nil + + case gitInfoMsg: + m.statusBar.SetGit(msg) + return m, nil + + case tea.BackgroundColorMsg: + // Feed the terminal's real background to the Markdown renderer so glamour + // picks a matching light/dark palette WITHOUT issuing its own terminal + // query (which would leak its reply into the input — see SetMarkdownDark). + // Re-flow so any already-finalized assistant block re-renders in the right + // palette. + SetMarkdownDark(msg.IsDark()) + m.transcript.reflow() + return m, nil + + case tea.MouseWheelMsg: + // Mouse-wheel scrolling reaches the transcript viewport whether idle or + // running, so history stays scrollable with the wheel — not just PgUp/PgDn. + // The viewport (MouseWheelEnabled by default) turns the wheel event into a + // scroll; enabling MouseModeCellMotion in View is what makes the terminal + // deliver these events under the alt-screen at all. + cmd := m.transcript.update(msg) + m.sel = selection{} + return m, cmd + + case tea.MouseClickMsg: + // A left press on the scrollbar column grabs the thumb (jump + drag). A left + // press anywhere else begins a text selection at that cell, replacing any + // prior one; a bare click (no drag) leaves it empty so it clears the old + // highlight without starting a copyable range. + if msg.Button == tea.MouseLeft { + if m.onScrollbar(msg.X, msg.Y) { + m.draggingScrollbar = true + m.transcript.scrollToRow(msg.Y) + return m, nil + } + m.sel = selection{active: true, anchor: point{msg.X, msg.Y}, cursor: point{msg.X, msg.Y}} + return m, nil + } + return m, nil + + case tea.MouseMotionMsg: + // While the thumb is grabbed, vertical motion drags it regardless of the + // cursor's column. Otherwise, motion after a left press extends the text + // selection to the current cell. + if m.draggingScrollbar { + m.transcript.scrollToRow(msg.Y) + return m, nil + } + if m.sel.active { + m.sel.cursor = point{msg.X, msg.Y} + } + return m, nil + + case tea.MouseReleaseMsg: + m.draggingScrollbar = false + if m.sel.active { + m.sel.cursor = point{msg.X, msg.Y} + } + return m, nil + + case tea.PasteMsg: + // Bracketed paste (e.g. Cmd+V / right-click paste): the terminal delivers + // the whole clipboard payload as one message. A multi-line paste is + // collapsed to a compact placeholder (expanded at submit); a single-line + // paste is inserted verbatim. See handlePaste. + if !m.running { + return m.handlePaste(msg.Content) + } + return m, nil + + case tea.ClipboardMsg: + // OSC52 clipboard read reply (from tea.ReadClipboard on Ctrl+V / Cmd+V). + // Route through the same collapse-or-insert path as bracketed paste. + if !m.running { + return m.handlePaste(msg.Content) + } + return m, nil + + case clipboardImageMsg: + // Reply to a Ctrl+V / Cmd+V image-read attempt. With an image, drop an + // "[Image #N]" placeholder (expanded to an @image reference at submit); with + // none, fall back to a normal OSC52 text read so plain-text paste still works. + if !m.running { + if msg.ok { + return m.handleImagePaste(msg.path) + } + return m, tea.ReadClipboard + } + return m, nil + + case tea.KeyPressMsg: + return m.handleKey(msg) + + case spinnerTickMsg: + // Advance the working animation and schedule the next frame, but only while + // a run is in flight; once idle the tick is not re-issued so the spinner + // stops without a lingering goroutine. + if !m.running { + return m, nil + } + m.spinner.advance() + return m, m.tickSpinner() + + case textDeltaMsg: + m.spinner.addTokens(msg.delta) + m.transcript.appendDelta(msg.delta) + m.remoteEcho(msg.delta) + return m, m.pumpNext() + + case turnEndMsg: + m.transcript.finalizeTurn(msg.msg) + // Surface a failed or empty turn so a provider/API error is never silent. + // The loop delivers request failures (e.g. a 4xx from the endpoint) as a + // terminal assistant message with stopReason error/aborted via TurnEndEvent + // — not as the run's result error (runEndMsg.err) — so without this check + // the TUI would finalize an empty turn and return to the prompt with no + // output at all. Mirrors the headless driver and the line-based REPL. + switch msg.msg.StopReason { + case agentcore.StopReasonError: + reason := strings.TrimSpace(msg.msg.ErrorMessage) + if reason == "" { + reason = "the provider returned an error with no message" + } + m.transcript.addSystem("error: " + reason) + case agentcore.StopReasonAborted: + m.transcript.addSystem("error: aborted") + default: + // A turn that ends cleanly (end_turn) but produced no text, no thinking, + // and no tool calls means the endpoint accepted the request but sent back + // nothing usable (e.g. a 200 whose body was not in the wire format this + // protocol expects). Note it instead of showing nothing. + if len(msg.msg.Content) == 0 && len(msg.results) == 0 { + m.transcript.addSystem("note: empty response from the model (no content). " + + "Check that --model, --base-url and --protocol match the same provider.") + } + } + return m, m.pumpNext() + + case toolStartMsg: + // Create a rich tool-call card, index it by id for the later end event, and + // append it as an ordered transcript block so it renders inline (#389). + card := &toolCard{id: msg.id, name: msg.name, input: msg.input, state: cardRunning} + m.toolCards[msg.id] = card + m.lastToolCard = card + m.transcript.addToolCard(card) + m.remoteEcho("\n· " + msg.name + "\n") + // A `task` tool call dispatches a sub-agent: open a status-panel row keyed by + // the tool-call id (matching the later progress/end events) and record its + // start so elapsed can be shown live (SPEC 4.4). + if msg.name == "task" { + m.subagents.add(msg.id, taskDescription(msg.input), time.Now()) + m.relayout() // the new panel row shrinks the transcript to fit + } + return m, m.pumpNext() + + case toolUpdateMsg: + // A `task` sub-agent forwards its text as incremental tool-update deltas; + // accumulate them onto the matching panel row so the expanded view can show + // the running output. appendOutput is a no-op for non-task ids (nothing to + // attach to), so ordinary tool updates are unaffected. Relayout only when the + // delta lands on the currently expanded row, whose growing output changes the + // panel height; other rows' output is buffered without touching the layout. + m.subagents.appendOutput(msg.id, msg.partial) + if m.subagents.expandedID() == msg.id { + m.relayout() + } + return m, m.pumpNext() + + case subagentProgressMsg: + // A running sub-agent reported structured progress: refresh its panel row's + // activity/tokens. update adds the row if it is missing so a late/out-of-order + // progress (arriving before the task's start) is still shown (SPEC 5.4). + m.subagents.update(msg.id, msg.desc, msg.activity, msg.tokens, time.Now()) + m.relayout() // a first-seen id adds a row; keep the transcript sized to it + return m, m.pumpNext() + + case toolEndMsg: + // Flip the card's state and attach the parsed response tree. The card is + // held by pointer in the transcript, so a reflow re-renders it in place. + if card, ok := m.toolCards[msg.id]; ok { + if msg.ok { + card.state = cardSuccess + } else { + card.state = cardWarn + } + card.response = parseToolResult(msg.result) + m.transcript.reflow() + } + // Retire the sub-agent's status-panel row (a no-op for non-task tools whose id + // was never added), reclaiming its reserved height. + if _, wasSub := m.subagents.byID[msg.id]; wasSub { + m.subagents.remove(msg.id) + m.relayout() + } + return m, m.pumpNext() + + case telemetryMsg: + // Feed the status bar's context-usage readout, and retain the event on the + // session's telemetry holder so /status can render the cumulative + last-run + // telemetry report (US-002, #292). Then keep the pump running. + m.statusBar.SetTelemetry(telemetryEventView{ + util: msg.ev.ContextUtilization, + window: msg.ev.ContextWindow, + tokens: msg.ev.ContextTokens, + }) + if m.session != nil && m.session.telemetry != nil { + m.session.telemetry.Fold(msg.ev) + } + return m, m.pumpNext() + + case compactionStartMsg: + m.spinner.pin("Compacting conversation") + return m, m.pumpNext() + + case compactionMsg: + m.spinner.unpin() + if m.session != nil { + m.session.compacted = true + } + m.transcript.addSystem("(context compacted)") + return m, m.pumpNext() + + case rebuildDoneMsg: + // A manual /rebuild finished: clear the pinned "Preparing conversation + // context…" spinner (no run is pumping, so stop it and drop out of the + // running state) and report the outcome. rebuild() already applied the + // rebuilt messages and set session.compacted on success. + m.spinner.unpin() + m.spinner.stop() + m.running = false + if msg.err != nil { + m.transcript.addSystem("rebuild failed: " + msg.err.Error() + " (context left unchanged)") + } else { + m.transcript.addSystem(msg.summary) + } + m.relayout() + return m, nil + + case runEndMsg: + m.running = false + m.runCh = nil + m.spinner.stop() + // The run is over: any still-open sub-agent rows are stale (their tasks ended + // with the run), so clear the panel to reclaim its height. + m.subagents = subagentPanel{} + m.relayout() + if msg.err != nil { + m.transcript.addSystem("Run ended: " + msg.err.Error()) + } + // Persist the turn's new messages as a branch so the conversation survives + // exit and can be resumed (FR-16). This is race-free: the pump goroutine + // owns agentCtx.Messages during the run and only sends runEndMsg after + // DrainStream returns (loop done), so no goroutine is still writing the + // context when persist reads it here on the tea goroutine. A save failure + // is surfaced but non-fatal. + if m.session != nil { + if err := m.session.persist(); err != nil { + m.transcript.addSystem("Session save failed: " + err.Error()) + } + } + // The editor was blurred at submit; re-enable it so the next prompt can be + // typed, and re-probe git since a run may have changed the working tree. + focus := m.input.Focus() + return m, tea.Batch(focus, fetchGitCmd(m.cwd)) + + case remoteInputMsg: + // A prompt arrived from the paired browser (remote-control). Always re-issue + // the listener so successive remote prompts keep arriving. While a run is in + // flight the prompt is refused with a note (mirroring the local single-run + // gate); when idle it is echoed as a user block and run — as a slash command + // if it starts with "/", else a normal prompt. + text := strings.TrimSpace(msg.text) + if m.running || text == "" { + if m.running && text != "" { + m.transcript.addSystem("(remote input ignored: a run is in progress)") + m.relayout() + } + return m, m.waitRemoteInput() + } + var cmd tea.Cmd + var next tea.Model = m + if strings.HasPrefix(text, "/") { + next, cmd = m.runSlash(text) + } else { + m.transcript.addUser(text) + m.remoteEcho("\n> " + text + "\n") + m.relayout() + next, cmd = m.startPrompt(text) + } + m = next.(Model) + return m, tea.Batch(cmd, m.waitRemoteInput()) + } + return m, nil +} + +// handleKey processes a key press. It resolves the keys the shell owns — +// two-stage interrupt/quit, prompt submit, transcript scrolling — and delegates +// everything else (character entry, in-buffer cursor movement, Shift+Enter +// newline) to the input editor while idle. Keys are matched via KeyPressMsg +// .String() so the mapping is terminal-independent. +func (m Model) handleKey(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + // While idle with the autocomplete popup open, the arrow / Tab / Esc keys + // drive the menu instead of the transcript or textarea (FR-15). Enter is left + // to the main switch below, which routes through submit → runSlash so the + // selected/typed command runs. These are matched via KeyPressMsg.String() so + // the mapping is terminal-independent. + if !m.running && m.menu.active { + switch msg.String() { + case "up": + m.menu.moveUp() + return m, nil + case "down": + m.menu.moveDown() + return m, nil + case "tab": + m = m.completeSlash() + m.relayout() + return m, nil + case "esc": + m.menu.close() + m.relayout() + return m, nil + case "enter": + return m.submitSlashSelected() + } + } + + // While a sub-agent run is streaming, the composer is disabled (no typing until + // the run ends), so ↑/↓ drive a selection cursor over the live sub-agent status + // rows and Enter expands the selected row to show its accumulated output inline. + // Esc is the one-key escape back to the composer: with a row selected it drops + // the selection AND re-focuses the input box in a single press, so arrowing into + // the panel is never a trap. With no selection Esc falls through to its + // two-stage interrupt role below. The Value()=="" guard is a safety net for the + // rare case where text reached the buffer (e.g. a paste): then arrows edit the + // buffer rather than the panel. + if m.running && m.subagents.active() > 0 && m.input.Value() == "" { + switch msg.String() { + case "up": + m.subagents.selectUp() + m.relayout() + return m, nil + case "down": + m.subagents.selectDown() + m.relayout() + return m, nil + case "enter": + m.subagents.toggleExpand() + m.relayout() + return m, nil + case "esc": + if m.subagents.hasSelection() { + m.subagents.clearSelection() + focus := m.input.Focus() + m.relayout() + return m, focus + } + } + } + + switch msg.String() { + case "ctrl+c": + // Ctrl+C copies the current mouse selection when there is one (over OSC52), + // clearing it afterward; with no selection it keeps its interrupt-or-quit + // role. Copying works even mid-run, so grabbing streamed output never + // interrupts the run. + if !m.sel.empty() { + text := m.selectedText() + m.sel = selection{} + if text != "" { + return m, tea.SetClipboard(text) + } + return m, nil + } + return m.interruptOrQuit() + case "super+c": + // Cmd+C on macOS is the platform-standard copy: copy the mouse selection + // when there is one (clearing it), else the whole input buffer. Unlike + // Ctrl+C it never interrupts/quits — Cmd+C means "copy" on macOS. Most + // terminals intercept Cmd+C for their own native copy and never deliver it + // here; this branch serves terminals that forward the Super modifier. + if !m.sel.empty() { + text := m.selectedText() + m.sel = selection{} + if text != "" { + return m, tea.SetClipboard(text) + } + return m, nil + } + if !m.running { + if v := m.input.Value(); v != "" { + return m, tea.SetClipboard(v) + } + } + return m, nil + case "esc": + return m.interruptOrQuit() + case "ctrl+o": + // Toggle the most-recent tool card between its capped preview and the full + // response tree, then re-flow so the change shows inline (#389). + if m.lastToolCard != nil { + m.lastToolCard.expanded = !m.lastToolCard.expanded + m.transcript.reflow() + } + return m, nil + case "ctrl+d": + // Ctrl+D quits only when idle; mid-run it is ignored so a run is never + // dropped by a stray EOF key. + if !m.running { + m.shutdownRemote() + m.quitting = true + return m, tea.Quit + } + return m, nil + case "enter": + // Enter submits the composed buffer (FR-13). Shift+Enter inserts a newline + // (rebound in newInput) so the editor is a true multi-line composer; when + // the slash menu is open, Enter runs the highlighted command (handled + // above), so this branch is only reached with the menu closed. + if !m.running { + return m.submit() + } + return m, nil + case "pgup", "pgdown": + // Page scrolling reaches the transcript viewport whether idle or running, + // so history stays readable while a run streams. Scrolling shifts the + // content under a screen-anchored selection, so drop the selection to avoid + // a stale highlight. Line-oriented keys (up / down / home / end) belong to + // the multi-line editor and are delegated below. + m.sel = selection{} + cmd := m.transcript.update(msg) + return m, cmd + case "ctrl+v": + // Explicit paste key: first try to pull an image off the clipboard (Claude + // Code-style image paste); the reply arrives as clipboardImageMsg and, when + // no image is present, falls back to an OSC52 text read (tea.ClipboardMsg). + // This is intercepted before textarea so its own Ctrl+V binding — which reads + // via an external process and returns an unexported message the model can't + // route — is bypassed. The common Cmd+V path does not reach here; it arrives + // as a bracketed tea.PasteMsg handled in Update. + if !m.running { + return m, readClipboardImage + } + return m, nil + case "super+v": + // Cmd+V on macOS is the platform-standard paste. Most terminals turn it + // into a bracketed paste (tea.PasteMsg, handled in Update); this branch + // covers terminals that instead forward the Super modifier as a key. Try an + // image read first, falling back to an OSC52 text read when none is present. + if !m.running { + return m, readClipboardImage + } + return m, nil + case "ctrl+y": + // Copy: the editor has no text selection, so this copies the whole buffer + // to the system clipboard over OSC52. A no-op on an empty buffer. + if !m.running { + if v := m.input.Value(); v != "" { + return m, tea.SetClipboard(v) + } + } + return m, nil + } + + // Everything else is editing input; gated on idle so keystrokes never corrupt + // an in-flight prompt. textarea handles CJK / emoji by rune and Shift+Enter as + // a newline. After the buffer changes, refresh the autocomplete popup so it + // opens/filters/closes as the user types a "/name" prefix. + if !m.running { + // ↑/↓ walk the submitted-prompt history when the caret is at the top / bottom + // edge of the composer; otherwise they move the caret within a multi-line + // draft (handled by the textarea below). + switch msg.String() { + case "up": + return m.historyPrev(msg) + case "down": + return m.historyNext(msg) + } + var cmd tea.Cmd + m.input, cmd = m.input.Update(msg) + m.menu.refresh(m.input.Value(), m.slash) + m.relayout() + return m, cmd + } + return m, nil +} + +// submit starts a run for the current buffer: it appends the user block, clears +// and blurs the editor, flips to running, and — when a run starter is wired — +// returns the first pump Cmd. With no starter (pre-#392) it records the prompt +// and a system note without launching anything, and leaves the editor ready for +// the next line. +func (m Model) submit() (tea.Model, tea.Cmd) { + raw := strings.TrimSpace(m.input.Value()) + prompt := strings.TrimSpace(m.expandImages(m.expandPastes(m.input.Value()))) + if prompt == "" { + return m, nil + } + // Record the input (as typed) into the browse history, then exit browse mode. + m.recordHistory(raw) + // The placeholders have been expanded into the prompt, so the stored paste + // bodies and image paths are consumed; drop them (the id counters keep climbing). + m.pastes = make(map[int]string) + m.images = make(map[int]string) + // A "/name ..." line is a slash-command invocation, not a prompt: resolve it + // against the shared registry (same as the REPL) rather than sending it to the + // agent verbatim. + if strings.HasPrefix(prompt, "/") { + return m.runSlash(prompt) + } + m.transcript.addUser(prompt) + m.remoteEcho("\n> " + prompt + "\n") + m.input.Clear() + m.menu.close() + m.relayout() + return m.startPrompt(prompt) +} + +// completeSlash fills the buffer with the highlighted candidate's "/name " so the +// user can go on to type arguments; the trailing space ends name-completion, so +// the refresh closes the popup. It is the Tab action while the menu is open. +func (m Model) completeSlash() Model { + if c, ok := m.menu.current(); ok { + m.input.SetValue("/" + c.Name + " ") + m.menu.refresh(m.input.Value(), m.slash) + } + return m +} + +// submitSlashSelected runs the command the popup highlights (Enter while the +// menu is open). Navigating with the arrows then pressing Enter runs the +// selected command even if the typed prefix is shorter; with no selection it +// falls back to the raw buffer so a fully-typed "/name" still runs. +func (m Model) submitSlashSelected() (tea.Model, tea.Cmd) { + line := strings.TrimSpace(m.input.Value()) + if c, ok := m.menu.current(); ok { + line = "/" + c.Name + } + m.recordHistory(line) + return m.runSlash(line) +} + +// runSlash resolves a slash-command line against the shared registry and folds +// its outcome into the transcript, mirroring the REPL's dispatch: the invocation +// is echoed as a user block; an action command's status (e.g. /help, /model) +// renders as a system block; a prompt/skill command's expanded text starts a +// run; a hybrid (plugin) command shows its notifications then runs its prompt. +// An unknown command surfaces the resolver error as a system block. +func (m Model) runSlash(line string) (tea.Model, tea.Cmd) { + // /exit and /quit terminate the TUI, mirroring the REPL loop which intercepts + // them before slash resolution. They register only as no-op /help builtins, so + // without this the registry would resolve them to an empty action. + if line == "/exit" || line == "/quit" { + m.shutdownRemote() + m.quitting = true + return m, tea.Quit + } + // /memory is intercepted before registry resolution (like /rebuild): it + // prints the persistent-memory + infinite-context report, reading the live + // memory store, memory root, session id, and messages that a slash Action + // closure (string→string) cannot reach. + if line == "/memory" || strings.HasPrefix(line, "/memory ") { + m.transcript.addUser(line) + m.input.Clear() + m.menu.close() + var buf bytes.Buffer + var store *memory.Store + var memoryRoot, sessionID string + var msgs agentcore.MessageList + window := m.live.ContextWindow + if m.session != nil { + store = m.session.memstore + memoryRoot = m.session.memoryRoot + sessionID = m.session.header.ID + msgs = m.session.agentCtx.Messages + } + memstatus.RunMemory(&buf, store, memoryRoot, sessionID, msgs, window) + m.transcript.addSystem(strings.TrimRight(buf.String(), "\n")) + m.relayout() + return m, nil + } + // /status is intercepted before registry resolution (like /memory): it prints + // the shared runtime/context/project/credentials/telemetry report, which reads + // the session's live collaborators (live config, trust manager, telemetry + // holder, slash registry) that a slash Action closure (string→string) cannot + // reach. The rendering lives in the shared status package so the TUI and the + // REPL produce byte-identical output. + if line == "/status" || strings.HasPrefix(line, "/status ") { + m.transcript.addUser(line) + m.input.Clear() + m.menu.close() + if m.session == nil { + m.transcript.addSystem("(status unavailable: no active session)") + m.relayout() + return m, nil + } + var buf bytes.Buffer + status.RunStatus(&buf, m.session) + m.transcript.addSystem(strings.TrimRight(buf.String(), "\n")) + m.relayout() + return m, nil + } + // /session is intercepted before registry resolution (like /memory): it prints + // the conversation summary (session id, message count, estimated tokens, model/ + // provider, created time, compaction count) from the live context — state a + // slash Action closure cannot reach. The rendering is shared with the REPL. + if line == "/session" { + m.transcript.addUser(line) + m.input.Clear() + m.menu.close() + if m.session == nil { + m.transcript.addSystem("(session unavailable: no active session)") + m.relayout() + return m, nil + } + var buf bytes.Buffer + m.session.renderSession(&buf) + m.transcript.addSystem(strings.TrimRight(buf.String(), "\n")) + m.relayout() + return m, nil + } + // /rebuild is intercepted before registry resolution (like /exit): it + // reconstructs the shared context from a persisted checkpoint (or falls back + // to compaction) and replaces the message list in place — work a slash Action + // closure cannot do. It reuses the compacting-indicator: the spinner is armed + // and pinned to "Preparing conversation context…" while the rebuild runs off + // the tea loop, and rebuildDoneMsg clears it and reports the result. + if line == "/rebuild" { + m.transcript.addUser(line) + m.input.Clear() + m.menu.close() + if m.session == nil { + m.transcript.addSystem("(rebuild unavailable: no active session)") + m.relayout() + return m, nil + } + m.spinner.begin(time.Now(), m.thinkingLabel()) + m.spinner.pin("Preparing conversation context") + m.running = true + m.relayout() + return m, tea.Batch(m.session.rebuildCmd(), m.tickSpinner()) + } + // /remote-control is intercepted before registry resolution (like /rebuild): + // it starts/stops the LAN mirror server, which owns state (server, bridge, + // listener Cmd) a string→string slash Action cannot hold. + if line == "/remote-control" || strings.HasPrefix(line, "/remote-control ") { + return m.runRemoteControl(line) + } + m.transcript.addUser(line) + m.input.Clear() + m.menu.close() + m.relayout() + if m.slash == nil { + m.transcript.addSystem("Slash commands unavailable") + return m, nil + } + outcome, err := m.slash.ResolveOutcome(line) + if err != nil { + m.transcript.addSystem(err.Error()) + return m, nil + } + if outcome.Message != "" { + m.transcript.addSystem(outcome.Message) + } + // A live-state command (/model, /think) may have mutated m.live; sync the + // status bar so the model/thinking segments reflect the switch immediately. + if m.live != nil { + m.statusBar.SetModel(m.live.Model) + m.statusBar.SetThinking(string(m.live.ThinkingLevel)) + } + // An action command is complete once its status is shown; a hybrid with no + // prompt (notifications only) likewise starts no run. + if outcome.Kind == runtime.SlashAction || outcome.Prompt == "" { + return m, nil + } + return m.startPrompt(outcome.Prompt) +} + +// recordHistory appends an submitted input to the browse history (skipping a +// consecutive duplicate, like a shell) and resets the browse cursor to the live +// draft, so the next ↑ starts from the most recent entry and any stashed draft is +// dropped. A blank entry is never stored. +func (m *Model) recordHistory(entry string) { + entry = strings.TrimSpace(entry) + if entry != "" && (len(m.history) == 0 || m.history[len(m.history)-1] != entry) { + m.history = append(m.history, entry) + } + m.histIdx = len(m.history) + m.histDraft = "" +} + +// historyPrev recalls the previous submitted input into the composer, but only +// when the caret is on the first line — otherwise ↑ moves the caret within a +// multi-line draft. The first recall stashes the live draft so historyNext can +// restore it, and the cursor lands past the newest entry (len(history)) initially. +func (m Model) historyPrev(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + if len(m.history) == 0 || m.input.Line() != 0 { + var cmd tea.Cmd + m.input, cmd = m.input.Update(msg) + m.menu.refresh(m.input.Value(), m.slash) + m.relayout() + return m, cmd + } + if m.histIdx == len(m.history) { + m.histDraft = m.input.Value() + } + if m.histIdx > 0 { + m.histIdx-- + } + m.input.SetValue(m.history[m.histIdx]) + m.menu.refresh(m.input.Value(), m.slash) + m.relayout() + return m, nil +} + +// historyNext walks forward toward more recent inputs — restoring the stashed +// draft once it steps past the newest entry — but only while browsing and with +// the caret on the last line; otherwise ↓ moves the caret within a multi-line +// draft. +func (m Model) historyNext(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) { + if m.histIdx >= len(m.history) || m.input.Line() != m.input.LineCount()-1 { + var cmd tea.Cmd + m.input, cmd = m.input.Update(msg) + m.menu.refresh(m.input.Value(), m.slash) + m.relayout() + return m, cmd + } + m.histIdx++ + if m.histIdx == len(m.history) { + m.input.SetValue(m.histDraft) + } else { + m.input.SetValue(m.history[m.histIdx]) + } + m.menu.refresh(m.input.Value(), m.slash) + m.relayout() + return m, nil +} + +// startPrompt launches an agent run for prompt, blurring the editor and flipping +// to running when a run starter is wired. With no starter (pre-session model / +// tests) it records the pre-#392 system note and stays idle. It is shared by a +// plain submit and by a slash prompt/skill command. +func (m Model) startPrompt(prompt string) (tea.Model, tea.Cmd) { + if m.startRunFn == nil { + m.transcript.addSystem("(run not wired up: see session assembly in #392)") + return m, nil + } + m.input.Blur() + ch, cmd := m.startRunFn(prompt) + m.runCh = ch + m.running = true + m.spinner.begin(time.Now(), m.thinkingLabel()) + m.relayout() + return m, tea.Batch(cmd, m.tickSpinner()) +} + +// thinkingLabel returns the current thinking-effort label for the spinner stats +// (e.g. "medium"), or "" when no thinking level is configured so the stat is +// omitted. It reads the live config the /model command mutates, falling back to +// the launch Options. +func (m Model) thinkingLabel() string { + if m.live != nil && m.live.ThinkingLevel != "" { + return string(m.live.ThinkingLevel) + } + return string(m.opts.ThinkingLevel) +} + +// taskDescription pulls the human-readable "description" out of a `task` tool +// call's decoded arguments for the sub-agent panel's row label. It returns "" +// when absent or non-string (the description field is optional in the schema), +// in which case the panel row leads with the activity instead. +func taskDescription(input map[string]any) string { + if s, ok := input["description"].(string); ok { + return s + } + return "" +} + +// tickSpinner schedules the next spinner animation frame. The model re-issues it +// on each spinnerTickMsg while running, so the animation self-sustains until the +// run ends (the tick is simply not re-issued once idle). +func (m Model) tickSpinner() tea.Cmd { + return tea.Tick(spinnerInterval, func(t time.Time) tea.Msg { + return spinnerTickMsg(t) + }) +} + +// interruptOrQuit is the shared Esc / bare-Ctrl+C action: a two-stage interrupt +// (FR-14) that stops an in-flight run on the first press and stays in the +// program, or quits when idle. +func (m Model) interruptOrQuit() (tea.Model, tea.Cmd) { + if m.running { + if m.interruptFn != nil { + m.interruptFn() + } + m.transcript.addSystem("(interrupting the current run…)") + return m, nil + } + m.shutdownRemote() + m.quitting = true + return m, tea.Quit +} + +// shutdownRemote stops the remote-control server on quit so the listener and +// WebSocket are released cleanly. A no-op when remote control is off or no +// session is bound. +func (m Model) shutdownRemote() { + if m.session != nil { + m.session.stopRemote() + } +} + +// feedInput forwards a message (a paste payload) to the editor, then refreshes +// the slash menu and re-lays out because inserted text can add lines (growing +// the editor) or begin a "/name". It is the shared tail of the paste handlers. +func (m Model) feedInput(msg tea.Msg) (tea.Model, tea.Cmd) { + var cmd tea.Cmd + m.input, cmd = m.input.Update(msg) + m.menu.refresh(m.input.Value(), m.slash) + m.relayout() + return m, cmd +} + +// pastePlaceholderRe matches the "[Pasted text #N +M lines]" tokens handlePaste +// leaves in the composer, capturing the id so expandPastes can swap the stored +// body back in at submit. +var pastePlaceholderRe = regexp.MustCompile(`\[Pasted text #(\d+) \+\d+ lines\]`) + +// handlePaste inserts a pasted payload into the editor. A multi-line paste is +// collapsed to a compact "[Pasted text #N +M lines]" placeholder (the full body +// stashed in m.pastes for expansion at submit), so a large paste does not flood +// the composer — mirroring Claude Code. A single-line paste is inserted verbatim. +func (m Model) handlePaste(content string) (tea.Model, tea.Cmd) { + if content == "" { + return m, nil + } + if strings.Contains(content, "\n") { + if m.pastes == nil { + m.pastes = make(map[int]string) + } + m.pasteSeq++ + id := m.pasteSeq + m.pastes[id] = content + lines := strings.Count(content, "\n") + 1 + placeholder := fmt.Sprintf("[Pasted text #%d +%d lines]", id, lines) + return m.feedInput(tea.PasteMsg{Content: placeholder}) + } + return m.feedInput(tea.PasteMsg{Content: content}) +} + +// expandPastes replaces every paste placeholder in s with its stored body, so +// the submitted prompt carries the real pasted text rather than the compact +// token the user saw in the composer. An unknown id (e.g. the user edited the +// token) is left as-is. It returns s unchanged when no pastes are stashed. +func (m Model) expandPastes(s string) string { + if len(m.pastes) == 0 { + return s + } + return pastePlaceholderRe.ReplaceAllStringFunc(s, func(tok string) string { + sm := pastePlaceholderRe.FindStringSubmatch(tok) + id, err := strconv.Atoi(sm[1]) + if err != nil { + return tok + } + if body, ok := m.pastes[id]; ok { + return body + } + return tok + }) +} + +// handleImagePaste stashes a pasted image (already saved to a temp PNG at path) +// and drops a compact "[Image #N]" placeholder into the composer, mirroring the +// text-paste placeholder. submit expands it into an "@image:" reference so +// BuildUserContent attaches the image as multimodal content. An empty path falls +// back to a plain text read. +func (m Model) handleImagePaste(path string) (tea.Model, tea.Cmd) { + if path == "" { + return m, tea.ReadClipboard + } + if m.images == nil { + m.images = make(map[int]string) + } + m.imageSeq++ + id := m.imageSeq + m.images[id] = path + placeholder := fmt.Sprintf("[Image #%d]", id) + return m.feedInput(tea.PasteMsg{Content: placeholder}) +} + +// imagePlaceholderRe matches the "[Image #N]" tokens handleImagePaste leaves in +// the composer, capturing the id so expandImages can swap the stored temp path +// back in as an "@image:" reference at submit. +var imagePlaceholderRe = regexp.MustCompile(`\[Image #(\d+)\]`) + +// expandImages replaces every image placeholder in s with an "@image:" +// reference so BuildUserContent reads and attaches the pasted image. An unknown id +// (e.g. the user edited the token) is left as-is. It returns s unchanged when no +// images are stashed. +func (m Model) expandImages(s string) string { + if len(m.images) == 0 { + return s + } + return imagePlaceholderRe.ReplaceAllStringFunc(s, func(tok string) string { + sm := imagePlaceholderRe.FindStringSubmatch(tok) + id, err := strconv.Atoi(sm[1]) + if err != nil { + return tok + } + if p, ok := m.images[id]; ok { + return "@image:" + p + } + return tok + }) +} + +// pumpNext re-issues waitForEvent for the in-flight run so the next bridged msg +// is pulled. It returns nil once the run has ended (runCh cleared), stopping the +// pump. +func (m Model) pumpNext() tea.Cmd { + if m.running && m.runCh != nil { + return waitForEvent(m.runCh) + } + return nil +} + +// View implements tea.Model. It renders the shell on the alt-screen: the +// scrolling transcript filling the top rows, then the autocomplete popup (when +// open) and the multi-line input editor, and finally the persistent status bar +// (#386) on the very bottom row — below the input, per the layout fix. Setting +// AltScreen on the returned View is how Bubble Tea v2 enters/leaves the alternate +// screen buffer, so the user's scrollback is restored on quit. +func (m Model) View() tea.View { + if m.quitting { + return tea.View{AltScreen: true} + } + + content := m.applySelection(m.renderContent()) + + // MouseModeCellMotion enables click/release/wheel events. Without it the + // alt-screen swallows the wheel (no native scrollback), so history could only + // be reached via PgUp/PgDn; enabling it lets the wheel scroll the transcript + // and drives both scrollbar drag and mouse text selection. + return tea.View{Content: content, AltScreen: true, MouseMode: tea.MouseModeCellMotion} +} + +// renderContent builds the full-screen shell string (transcript, autocomplete +// popup, input editor, status bar) without any selection overlay. View wraps it +// with applySelection for display, and selectedText reuses it to extract the +// copied text from the exact rows the user sees. +func (m Model) renderContent() string { + width := m.width + if width <= 0 { + width = 80 + } + height := m.height + if height <= 0 { + height = 24 + } + + status := m.statusBar.Render(width) + + // The input editor renders its own prompt column and cursor across as many + // rows as the buffer currently spans (up to maxInputRows). + input := m.input.View() + + // Fallback transcript rows before the first size message; once sized the + // viewport is pre-sized by relayout and pads its own content. + rows := transcriptHeight(height) + sized := m.width > 0 && m.height > 0 + + var b strings.Builder + if sized { + // The viewport pads its content to exactly the rows relayout reserved. + b.WriteString(m.transcript.view()) + b.WriteByte('\n') + } else { + for i := 0; i < rows; i++ { + b.WriteByte('\n') + } + } + // The working spinner sits on its own row just above the input while a run is + // in flight (relayout reserves the row so the transcript shrinks to fit). The + // sub-agent status panel, when any `task` sub-agents are live, renders on the + // rows just ABOVE the spinner: one line each, elapsed refreshed every tick. + if m.running { + if panel := m.subagents.view(m.theme, width, time.Now()); panel != "" { + b.WriteString(panel) + b.WriteByte('\n') + } + if line := m.spinner.view(width); line != "" { + b.WriteString(line) + b.WriteByte('\n') + } + } + // The autocomplete popup, when open, renders just above the input line as an + // overlay (it contributes no rows while idle, so the empty-shell layout is + // unchanged). + if menu := m.menu.view(width); menu != "" { + b.WriteString(menu) + b.WriteByte('\n') + } + b.WriteString(input) + b.WriteByte('\n') + // The status bar is the final line, pinned to the very bottom of the shell + // below the input editor. + b.WriteString(status) + return b.String() +} + +// applySelection overlays the mouse selection highlight onto the rendered +// content, inverting the selected cells like a terminal's own selection. Only +// rows the selection intersects are rewritten (as plain text with the span +// inverted); untouched rows keep their original coloring. It is a no-op when the +// selection is empty. +func (m Model) applySelection(content string) string { + if m.sel.empty() { + return content + } + start, end := m.sel.ordered() + hi := lipgloss.NewStyle().Reverse(true) + rows := strings.Split(content, "\n") + for y := start.y; y <= end.y && y < len(rows); y++ { + if y < 0 { + continue + } + c0, c1, ok := rowRange(start, end, y) + if !ok { + continue + } + rows[y], _ = selectRow(rows[y], c0, c1, hi) + } + return strings.Join(rows, "\n") +} + +// selectedText extracts the plain text under the current selection from the rows +// the user sees, joining rows with newlines and trimming each row's trailing +// padding so copied text has no ragged whitespace tail. It returns "" when the +// selection is empty. +func (m Model) selectedText() string { + if m.sel.empty() { + return "" + } + start, end := m.sel.ordered() + rows := strings.Split(m.renderContent(), "\n") + var b strings.Builder + wrote := false + for y := start.y; y <= end.y && y < len(rows); y++ { + if y < 0 { + continue + } + c0, c1, ok := rowRange(start, end, y) + if !ok { + continue + } + _, text := selectRow(rows[y], c0, c1, lipgloss.Style{}) + if wrote { + b.WriteByte('\n') + } + b.WriteString(strings.TrimRight(text, " ")) + wrote = true + } + return b.String() +} + +// relayout re-sizes the transcript to the rows left after reserving the status +// bar (1 row), the current input editor height, and any open autocomplete popup. +// It hands the transcript the full width; the transcript itself spends one column +// on the scrollbar only while its content overflows (see transcript.reflow), so a +// short conversation uses the whole width and shows no bar, while a scrolling one +// reserves the gutter — and that decision re-runs on every streamed line, not just +// on resize. It is called on every resize and after any edit that changes the +// input height or menu row count. +func (m *Model) relayout() { + if m.width <= 0 || m.height <= 0 { + return + } + rows := m.height - 1 - m.input.Height() - m.menu.rows() + if m.running { + rows-- // the working spinner occupies the row just above the input + // The sub-agent panel reserves one status row per live sub-agent, plus the + // wrapped output lines of the expanded row (if any); an empty panel reserves + // nothing so the single-run layout is unchanged. + rows -= m.subagents.lineCount(m.width) + } + if rows < 0 { + rows = 0 + } + m.transcript.setSize(m.width, rows) + m.input.SetWidth(m.width) +} + +// onScrollbar reports whether the terminal cell (x, y) is the transcript's +// scrollbar: the rightmost column (relayout reserves m.width-1 for content, so +// the bar sits at column m.width-1) within the transcript's visible rows, which +// start at the top of the screen (row 0). It gates click-to-drag so presses in +// the body or on other chrome are left alone. When the content fits there is no +// bar (relayout reclaims the column), so it always returns false. +func (m Model) onScrollbar(x, y int) bool { + if m.width <= 0 || !m.transcript.overflowing() { + return false + } + h := m.transcript.viewportHeight() + return x == m.width-1 && y >= 0 && y < h +} + +// transcriptHeight returns the fallback number of rows for the transcript before +// the first size message arrives: the total minus the status bar and a single +// input row, floored at zero so tiny terminals never produce a negative extent. +func transcriptHeight(total int) int { + h := total - 2 + if h < 0 { + h = 0 + } + return h +} diff --git a/pigo/internal/cli/tui/model_test.go b/pigo/internal/cli/tui/model_test.go new file mode 100644 index 0000000..197c01d --- /dev/null +++ b/pigo/internal/cli/tui/model_test.go @@ -0,0 +1,519 @@ +package tui + +import ( + "fmt" + "strings" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// TestModelQuitKeys verifies the root model returns tea.Quit on the standard +// exit keys (Ctrl+C / Ctrl+D), which is how Bubble Tea tears down the program +// and restores the terminal from the alt-screen. +func TestModelQuitKeys(t *testing.T) { + for _, key := range []string{"ctrl+c", "ctrl+d"} { + m := NewModel(Options{}) + got, cmd := m.Update(keyPress(key)) + if cmd == nil { + t.Fatalf("%s: expected a quit command, got nil", key) + } + if msg := cmd(); msg != (tea.QuitMsg{}) { + t.Errorf("%s: cmd produced %T, want tea.QuitMsg", key, msg) + } + if !got.(Model).quitting { + t.Errorf("%s: model should be marked quitting", key) + } + } +} + +// TestModelViewShell verifies the empty shell renders on the alt-screen and, +// once a size is known, occupies the full terminal height (empty transcript rows +// + status bar + input line), with the real status bar (#386) painting its +// fields. +func TestModelViewShell(t *testing.T) { + m := NewModel(Options{Model: "test-model"}) + next, _ := m.Update(tea.WindowSizeMsg{Width: 60, Height: 10}) + view := next.View() + if !view.AltScreen { + t.Error("View should request the alt-screen") + } + if got := strings.Count(view.Content, "\n"); got != 9 { + t.Errorf("newline count = %d, want 9 (10 rows)", got) + } + if !strings.Contains(view.Content, "test-model") { + t.Errorf("status bar model field missing from view: %q", view.Content) + } +} + +// TestModelNewlineKeys verifies that Shift+Enter inserts a line break at the +// cursor and preserves the already-typed text, rather than submitting. Plain +// Enter still submits, so it does not leave a newline in the buffer. Shift+Enter +// is the primary newline key (reported distinctly by terminals speaking the +// Kitty disambiguate protocol, which Bubble Tea enables by default); Ctrl+J and +// Alt+Enter are fallbacks for terminals that collapse Shift+Enter to a bare CR. +func TestModelNewlineKeys(t *testing.T) { + var mm tea.Model = NewModel(Options{}) + mm, _ = mm.Update(tea.WindowSizeMsg{Width: 60, Height: 10}) + for _, r := range "abc" { + mm, _ = mm.Update(tea.KeyPressMsg{Code: r, Text: string(r)}) + } + mm, _ = mm.Update(tea.KeyPressMsg{Code: tea.KeyEnter, Mod: tea.ModShift}) + for _, r := range "def" { + mm, _ = mm.Update(tea.KeyPressMsg{Code: r, Text: string(r)}) + } + if got := mm.(Model).input.Value(); got != "abc\ndef" { + t.Errorf("input = %q, want %q", got, "abc\ndef") + } +} + +// TestModelSelectionCopy drives a mouse selection over a transcript line and +// asserts Ctrl+C copies the selected text (over OSC52) and clears the selection. +func TestModelSelectionCopy(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + m.transcript.addUser("hello world") + + // Locate the rendered screen cell where the text begins so the test does not + // hard-code the transcript's bottom-stick row. + rows := strings.Split(m.renderContent(), "\n") + y, x := -1, -1 + for i, r := range rows { + plain := stripANSI(r) + if idx := strings.Index(plain, "hello world"); idx >= 0 { + y = i + x = ui.Width(plain[:idx]) + break + } + } + if y < 0 { + t.Fatal("rendered screen did not contain the transcript text") + } + + // Select exactly "hello world" (11 display cells) on that row. + m.sel = selection{active: true, anchor: point{x, y}, cursor: point{x + 11, y}} + next, cmd := m.Update(keyPress("ctrl+c")) + if cmd == nil { + t.Fatal("ctrl+c with a selection should emit a clipboard command") + } + if got := fmt.Sprintf("%s", cmd()); got != "hello world" { + t.Errorf("copied %q, want %q", got, "hello world") + } + if !next.(Model).sel.empty() { + t.Error("selection should be cleared after Ctrl+C copies it") + } +} + +// TestModelCtrlCFallsBackToQuit verifies Ctrl+C with no selection keeps its +// interrupt/quit role (idle → quit). +func TestModelCtrlCFallsBackToQuit(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + next, cmd := m.Update(keyPress("ctrl+c")) + if cmd == nil || cmd() != (tea.QuitMsg{}) { + t.Fatal("ctrl+c without a selection should quit when idle") + } + if !next.(Model).quitting { + t.Error("model should be marked quitting") + } +} + +// TestModelImagePasteInsertsPlaceholder verifies a clipboard image (already saved +// to a temp file) is stashed and shown in the composer as a compact "[Image #N]" +// placeholder, and that expandImages swaps it for an "@image:" reference at +// submit so BuildUserContent attaches it as multimodal content. +func TestModelImagePasteInsertsPlaceholder(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + next, _ := m.Update(clipboardImageMsg{path: "/tmp/pigo-clip-1.png", ok: true}) + m = next.(Model) + + if got, want := m.input.Value(), "[Image #1]"; got != want { + t.Errorf("composer showed %q, want placeholder %q", got, want) + } + if got := m.expandImages(m.input.Value()); got != "@image:/tmp/pigo-clip-1.png" { + t.Errorf("expandImages = %q, want the @image reference", got) + } +} + +// TestModelImagePasteFallsBackToText verifies an empty clipboard image reply +// (ok=false) falls back to an OSC52 text read rather than inserting anything. +func TestModelImagePasteFallsBackToText(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + next, cmd := m.Update(clipboardImageMsg{ok: false}) + if cmd == nil { + t.Fatal("no image on the clipboard should fall back to a text read command") + } + if got := next.(Model).input.Value(); got != "" { + t.Errorf("composer should stay empty on fallback, got %q", got) + } +} + +// TestModelExpandImagesUnknownID verifies an unknown image id is left untouched. +func TestModelExpandImagesUnknownID(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + m = apply(t, m, clipboardImageMsg{path: "/tmp/a.png", ok: true}) + + got := m.expandImages("see [Image #1] and [Image #7]") + want := "see @image:/tmp/a.png and [Image #7]" + if got != want { + t.Errorf("expandImages = %q, want %q", got, want) + } +} + +// keyPress builds a KeyPressMsg matching String()==s for the simple keys used +// in these tests (ctrl+). +func keyPress(s string) tea.KeyPressMsg { + switch s { + case "ctrl+c": + return tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl} + case "ctrl+d": + return tea.KeyPressMsg{Code: 'd', Mod: tea.ModCtrl} + case "ctrl+y": + return tea.KeyPressMsg{Code: 'y', Mod: tea.ModCtrl} + case "super+c": + return tea.KeyPressMsg{Code: 'c', Mod: tea.ModSuper} + case "super+v": + return tea.KeyPressMsg{Code: 'v', Mod: tea.ModSuper} + default: + return tea.KeyPressMsg{} + } +} + +// TestModelPasteSingleLineInsertsVerbatim verifies a single-line bracketed paste +// is inserted into the editor as-is (no placeholder collapsing). +func TestModelPasteSingleLineInsertsVerbatim(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + m = apply(t, m, tea.PasteMsg{Content: "hello world"}) + if got := m.input.Value(); got != "hello world" { + t.Errorf("input after paste = %q, want %q", got, "hello world") + } +} + +// TestModelPasteMultilineCollapses verifies a multi-line paste is collapsed to a +// compact "[Pasted text #N +M lines]" placeholder in the composer (Claude Code +// style) while the full body is stashed and expanded back at submit. +func TestModelPasteMultilineCollapses(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + m = apply(t, m, tea.PasteMsg{Content: "line1\nline2\nline3"}) + + if got, want := m.input.Value(), "[Pasted text #1 +3 lines]"; got != want { + t.Errorf("composer showed %q, want placeholder %q", got, want) + } + if got := m.expandPastes(m.input.Value()); got != "line1\nline2\nline3" { + t.Errorf("expandPastes = %q, want the original body", got) + } +} + +// TestModelExpandPastesMultiple verifies several collapsed pastes each expand +// back to their own body, and an unknown id is left untouched. +func TestModelExpandPastesMultiple(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + m = apply(t, m, tea.PasteMsg{Content: "aaa\nbbb"}) + m = apply(t, m, tea.PasteMsg{Content: "ccc\nddd"}) + + got := m.expandPastes("x [Pasted text #1 +2 lines] y [Pasted text #2 +2 lines] [Pasted text #9 +9 lines]") + want := "x aaa\nbbb y ccc\nddd [Pasted text #9 +9 lines]" + if got != want { + t.Errorf("expandPastes = %q, want %q", got, want) + } +} + +// TestModelClipboardReadInsertsIntoInput verifies an OSC52 clipboard read reply +// (tea.ClipboardMsg, the response to Ctrl+V) is inserted into the editor. +func TestModelClipboardReadInsertsIntoInput(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + m = apply(t, m, tea.ClipboardMsg{Content: "pasted"}) + if got := m.input.Value(); got != "pasted" { + t.Errorf("input after clipboard read = %q, want %q", got, "pasted") + } +} + +// TestModelCopyToClipboard verifies Ctrl+Y emits an OSC52 SetClipboard command +// carrying the current buffer, and is a no-op on an empty buffer. +func TestModelCopyToClipboard(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + // Empty buffer: no command. + if _, cmd := m.Update(keyPress("ctrl+y")); cmd != nil { + t.Errorf("ctrl+y on empty buffer should be a no-op, got a command") + } + + m = apply(t, m, tea.PasteMsg{Content: "copy me"}) + _, cmd := m.Update(keyPress("ctrl+y")) + if cmd == nil { + t.Fatal("ctrl+y with content should emit a clipboard command") + } + // SetClipboard yields an unexported string-underlying message; format it to + // read its payload without depending on the tea-internal type. + if got := fmt.Sprintf("%s", cmd()); got != "copy me" { + t.Errorf("clipboard command carried %q, want %q", got, "copy me") + } +} + +// TestModelSuperCCopiesSelection verifies Cmd+C (super+c) copies the mouse +// selection just like Ctrl+C, but with an empty buffer and no selection it is a +// no-op rather than quitting — Cmd+C is "copy" on macOS, never interrupt/quit. +func TestModelSuperCCopiesSelection(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + // No selection, empty buffer: no-op, and never a quit command. + if next, cmd := m.Update(keyPress("super+c")); cmd != nil { + t.Errorf("super+c with nothing to copy should be a no-op, got a command") + } else if next.(Model).quitting { + t.Error("super+c must never quit") + } + + // No selection, non-empty buffer: copies the whole buffer. + m = apply(t, m, tea.PasteMsg{Content: "buffer text"}) + if _, cmd := m.Update(keyPress("super+c")); cmd == nil { + t.Fatal("super+c with buffer content should emit a clipboard command") + } else if got := fmt.Sprintf("%s", cmd()); got != "buffer text" { + t.Errorf("super+c copied %q, want %q", got, "buffer text") + } +} + +// TestModelSuperVPastes verifies Cmd+V (super+v) requests the clipboard over +// OSC52 when idle, like Ctrl+V. +func TestModelSuperVPastes(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + if _, cmd := m.Update(keyPress("super+v")); cmd == nil { + t.Fatal("super+v should emit a clipboard read command when idle") + } +} + +// TestModelSubagentPanelLifecycle drives the sub-agent status panel through a +// task tool's lifecycle on the running model: a toolStartMsg(name=="task") opens +// a row, subagentProgressMsg refreshes it and it appears in the rendered View +// above the input, and the task's toolEndMsg retires it (empty panel → no rows). +func TestModelSubagentPanelLifecycle(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 20}) + m.running = true // the panel only renders while a run is in flight + m.spinner.begin(time.Now(), "") + + m = apply(t, m, toolStartMsg{id: "task-1", name: "task", input: map[string]any{"description": "build parser"}}) + if got := m.subagents.active(); got != 1 { + t.Fatalf("active after task start = %d, want 1", got) + } + + m = apply(t, m, subagentProgressMsg{id: "task-1", desc: "build parser", activity: "Editing", tokens: 64}) + if row := m.subagents.byID["task-1"]; row == nil || row.activity != "Editing" { + t.Fatalf("row after progress = %+v, want activity=Editing", row) + } + // The panel line is identified by its ⏺ glyph (distinct from the tool card, + // which also mentions the description) plus the live activity. + if view := m.View().Content; !strings.Contains(view, "⏺") || !strings.Contains(view, "Editing") { + t.Errorf("view missing panel line: %q", view) + } + + // A non-task tool must not open a panel row. + m = apply(t, m, toolStartMsg{id: "read-1", name: "read_file", input: map[string]any{"path": "/x"}}) + if got := m.subagents.active(); got != 1 { + t.Errorf("active after non-task start = %d, want 1", got) + } + + m = apply(t, m, toolEndMsg{id: "task-1", ok: true, result: "done"}) + if got := m.subagents.active(); got != 0 { + t.Errorf("active after task end = %d, want 0", got) + } + if view := m.View().Content; strings.Contains(view, "⏺") { + t.Errorf("view still shows retired panel line: %q", view) + } +} + +// TestModelCompactionIndicator verifies compactionStartMsg pins the spinner to +// "Compacting conversation…" while summarization runs, and compactionMsg clears +// the label and records the "(context compacted)" system note. +func TestModelCompactionIndicator(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 20}) + m.running = true + m.spinner.begin(time.Now(), "") + + m = apply(t, m, compactionStartMsg{}) + if view := stripANSI(m.spinner.view(120)); !strings.Contains(view, "Compacting conversation…") { + t.Errorf("spinner view %q should show the compaction label", view) + } + + m = apply(t, m, compactionMsg{}) + if m.spinner.pinned != "" { + t.Errorf("compactionMsg should unpin the spinner, got %q", m.spinner.pinned) + } + if joined := strings.Join(blockTexts(m.transcript), "\n"); !strings.Contains(joined, "(context compacted)") { + t.Errorf("transcript should note the compaction, got:\n%s", joined) + } +} + +// TestModelSubagentPanelHeightReservation verifies the panel's rows are reserved +// out of the transcript height so the total shell height is unchanged whether or +// not sub-agents are active. +func TestModelSubagentPanelHeightReservation(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 20}) + m.running = true + m.spinner.begin(time.Now(), "") + m.relayout() + base := m.transcript.viewportHeight() + + m = apply(t, m, toolStartMsg{id: "t1", name: "task", input: map[string]any{"description": "a"}}) + m = apply(t, m, toolStartMsg{id: "t2", name: "task", input: map[string]any{"description": "b"}}) + if got := m.transcript.viewportHeight(); got != base-2 { + t.Errorf("transcript height with 2 panel rows = %d, want %d (base %d - 2)", got, base-2, base) + } + // Every rendered frame stays exactly Height rows tall regardless of the panel. + if got := strings.Count(m.View().Content, "\n"); got != 19 { + t.Errorf("newline count = %d, want 19 (20 rows)", got) + } +} + +// TestModelSubagentPanelNavigation verifies that while a run streams and the +// composer is empty, ↓/↑ move the panel cursor, Enter expands the selected row's +// accumulated output inline (fed by tool-update deltas), and Esc collapses/clears +// the selection rather than interrupting the run. +func TestModelSubagentPanelNavigation(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 24}) + m.running = true + m.spinner.begin(time.Now(), "") + m = apply(t, m, toolStartMsg{id: "a", name: "task", input: map[string]any{"description": "task A"}}) + m = apply(t, m, toolStartMsg{id: "b", name: "task", input: map[string]any{"description": "task B"}}) + // A sub-agent's forwarded text arrives as an incremental tool-update delta. + m = apply(t, m, toolUpdateMsg{id: "b", partial: "output of B"}) + + // ↓ selects the top row, a second ↓ moves to row b. + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyDown}) + if !m.subagents.hasSelection() || m.subagents.selected != 0 { + t.Fatalf("after down: selected=%d hasSel=%v, want 0/true", m.subagents.selected, m.subagents.hasSelection()) + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyDown}) + + // Enter expands the selected row; its accumulated output shows in the render. + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyEnter}) + if got := m.subagents.expandedID(); got != "b" { + t.Fatalf("expandedID after enter = %q, want b", got) + } + if !strings.Contains(m.renderContent(), "output of B") { + t.Errorf("expanded render missing sub-agent output:\n%s", m.renderContent()) + } + + // Esc collapses/clears the selection and does NOT quit (no quit command). + next, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyEscape}) + m = next.(Model) + if m.subagents.hasSelection() { + t.Error("esc should clear the panel selection") + } + if cmd != nil { + if _, isQuit := cmd().(tea.QuitMsg); isQuit { + t.Error("esc with an active selection should not quit the program") + } + } +} + +// TestModelSubagentEscReturnsToInput verifies the one-key escape ("escape hatch: one key back to the input box"): +// while a sub-agent runs the composer is blurred (no typing), and after arrowing +// into the panel a single Esc both clears the selection and re-focuses the input +// box — so returning to the composer never requires more than one press and never +// interrupts the run. +func TestModelSubagentEscReturnsToInput(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 24}) + m.running = true + m.spinner.begin(time.Now(), "") + m.input.Blur() // the composer is blurred for the duration of a run (startPrompt) + m = apply(t, m, toolStartMsg{id: "a", name: "task", input: map[string]any{"description": "task A"}}) + + // Arrow into the panel: a selection is now active while the input stays blurred. + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyDown}) + if !m.subagents.hasSelection() { + t.Fatal("down should select a sub-agent row") + } + if m.input.Focused() { + t.Fatal("the composer should be blurred while a sub-agent run streams") + } + + // One Esc escapes: selection cleared AND the input box re-focused, in a single + // press, without interrupting the run. + next, _ := m.Update(tea.KeyPressMsg{Code: tea.KeyEscape}) + m = next.(Model) + if m.subagents.hasSelection() { + t.Error("one Esc should clear the panel selection") + } + if !m.input.Focused() { + t.Error("one Esc should re-focus the input box (return to the composer)") + } + if !m.running { + t.Error("escaping the panel selection must not interrupt the run") + } +} + +// TestModelPromptHistoryNavigation verifies shell-like prompt history: after +// submitting two prompts, ↑ from an empty composer recalls the most recent, a +// second ↑ walks further back, ↓ walks forward again, and a final ↓ restores the +// (empty) live draft. With no run starter wired, submit records history and +// leaves the composer idle/focused, so the arrow keys route to historyPrev / +// historyNext. +func TestModelPromptHistoryNavigation(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 12}) + + submit := func(s string) { + for _, r := range s { + m = apply(t, m, runeKey(r)) + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyEnter}) + } + submit("first prompt") + submit("second prompt") + + if m.input.Value() != "" { + t.Fatalf("composer should be empty after submit, got %q", m.input.Value()) + } + + // ↑ recalls the newest entry, a second ↑ the older one. + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyUp}) + if got := m.input.Value(); got != "second prompt" { + t.Errorf("first ↑ recalled %q, want %q", got, "second prompt") + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyUp}) + if got := m.input.Value(); got != "first prompt" { + t.Errorf("second ↑ recalled %q, want %q", got, "first prompt") + } + + // ↓ walks forward to the newer entry, then past it to restore the live draft. + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyDown}) + if got := m.input.Value(); got != "second prompt" { + t.Errorf("↓ walked to %q, want %q", got, "second prompt") + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyDown}) + if got := m.input.Value(); got != "" { + t.Errorf("↓ past newest should restore the empty draft, got %q", got) + } +} + +// TestModelPromptHistoryDedupsAndStashesDraft verifies two shell-like behaviors: +// a consecutive-duplicate submit is not stored twice, and an in-progress draft is +// stashed when browsing begins so ↓ past the newest entry brings it back. +func TestModelPromptHistoryDedupsAndStashesDraft(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 80, Height: 12}) + + submit := func(s string) { + for _, r := range s { + m = apply(t, m, runeKey(r)) + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyEnter}) + } + submit("same") + submit("same") // consecutive duplicate — must not be stored twice + if len(m.history) != 1 { + t.Fatalf("history = %v, want a single deduped entry", m.history) + } + + // Type a fresh draft, then browse: ↑ stashes the draft and recalls history, + // ↓ past the newest entry restores the stashed draft verbatim. + for _, r := range "draft" { + m = apply(t, m, runeKey(r)) + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyUp}) + if got := m.input.Value(); got != "same" { + t.Errorf("↑ recalled %q, want %q", got, "same") + } + m = apply(t, m, tea.KeyPressMsg{Code: tea.KeyDown}) + if got := m.input.Value(); got != "draft" { + t.Errorf("↓ should restore the stashed draft %q, got %q", "draft", got) + } +} diff --git a/pigo/internal/cli/tui/msgs.go b/pigo/internal/cli/tui/msgs.go new file mode 100644 index 0000000..fdd301f --- /dev/null +++ b/pigo/internal/cli/tui/msgs.go @@ -0,0 +1,82 @@ +package tui + +import "github.com/smallnest/pigo/internal/agentcore" + +// This file defines the tea.Msg types the event bridge (bridge.go) produces from +// a run's AgentEvents (US-004, SPEC 5.1). Each raw runtime signal is converted to +// exactly one of these value types so the Bubble Tea Update loop can dispatch on +// them with a plain type switch, keeping all run-time state changes on the tea +// goroutine (node #388 wires them into Model.Update). Every type is a value (not +// a pointer) so it flows through the tea.Msg (any) channel without aliasing the +// producer goroutine's state. + +// textDeltaMsg carries the newest suffix of streaming assistant text — the bytes +// produced since the previous delta for the current turn (see DrainStream's +// OnText contract). +type textDeltaMsg struct{ delta string } + +// turnEndMsg fires once per completed turn with the final assistant message and +// the tool results produced during it. +type turnEndMsg struct { + msg agentcore.AssistantMessage + results []agentcore.ToolResultMessage +} + +// toolStartMsg is emitted before a tool runs. input holds the decoded call +// arguments when they are a JSON object; it is nil otherwise (the raw Args are +// an untyped any at the event layer). +type toolStartMsg struct { + id string + name string + input map[string]any +} + +// toolUpdateMsg carries a partial result streamed during a tool's execution. +type toolUpdateMsg struct { + id string + partial string +} + +// toolEndMsg is emitted when a tool finishes. ok is false when the tool reported +// an error; result is the tool's textual output. +type toolEndMsg struct { + id string + ok bool + result string +} + +// subagentProgressMsg carries a running sub-agent's structured progress +// (translated from agentcore.SubAgentProgressEvent). id is the parent task +// tool-call id (the row key, matching the task's toolStartMsg/toolEndMsg id); +// desc is the task description (may be empty); activity is the current phase +// ("Reading"/"Editing"/…, never empty); tokens is a coarse output estimate +// (0 = unknown). Elapsed is NOT carried — the model computes it from the row's +// start time so the panel stays live without an event per frame. +type subagentProgressMsg struct { + id string + desc string + activity string + tokens int +} + +// telemetryMsg carries the run's end-of-run telemetry summary. +type telemetryMsg struct{ ev agentcore.TelemetryEvent } + +// compactionStartMsg signals that the loop is about to compact the context +// window. It pins the spinner to a "Compacting conversation…" label while the +// summarization request is in flight; compactionMsg clears it. +type compactionStartMsg struct{} + +// compactionMsg signals that the loop compacted the context window. The event's +// details are not needed by the transcript, so it is a bare signal. +type compactionMsg struct{} + +// runEndMsg is the final message: the run has fully drained. err is non-nil when +// the run ended in error (or was interrupted). +type runEndMsg struct{ err error } + +// remoteInputMsg carries a prompt submitted from the paired remote browser +// (remote-control, #443). The listener Cmd (Model.waitRemoteInput) blocks on the +// bridge's RemoteInput channel and emits one per submission, re-issued after each +// so successive remote prompts keep arriving. +type remoteInputMsg struct{ text string } diff --git a/pigo/internal/cli/tui/options.go b/pigo/internal/cli/tui/options.go new file mode 100644 index 0000000..0ac8332 --- /dev/null +++ b/pigo/internal/cli/tui/options.go @@ -0,0 +1,59 @@ +package tui + +import ( + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" +) + +// Options carries the resolved run configuration into Run. It deliberately +// mirrors repl.Options (see internal/cli/repl/interactive.go) field-for-field so +// cmd/pigo's dispatch can map the same assembled environment to either path, and +// so downstream nodes can port the REPL's session/live/slash/trust wiring into +// the TUI without reshaping the entry seam. This skeleton node does not yet +// consume most fields — they are here to lock the contract. +type Options struct { + Model string + ProviderName string + Provider provider.Provider + BaseURL string + APIKey string + Protocol string + // Version is the running build version (main.version), shown in the startup + // banner. Empty or "dev"/"unknown" renders as-is with no update hint. + Version string + // ThinkingLevel is the resolved reasoning-effort level (US-023): it seeds the + // live run config so every turn requests it, until a control command changes + // it. + ThinkingLevel agentcore.ThinkingLevel + Tools []agentcore.AgentTool + SysPrompt string + + // ResumeID, when non-empty, resumes an existing session: its messages seed + // the context and replayed transcript. Otherwise a fresh session is created. + ResumeID string + + // Approve, when true, grants the launch directory session trust before the + // run so the first-launch trust prompt is skipped and side-effect tools run + // without per-call confirmation (mirrors pi's --approve/-a). + Approve bool + // Skills is the pre-loaded skill set (loaded once by run.SetupEnv, shared with + // prompt injection). Each is registered as a /skill-name command. Empty under + // --no-skills, so nothing is registered. + Skills []*runtime.Skill + + // Plugins holds the loaded plugin manager so the TUI can deliver lifecycle + // events to subscribed plugins (US-017, #133). It may be nil (no plugins). + Plugins *plugin.Manager + + // ConfigPrompts holds prompt-template paths from the config.toml `prompts` + // array (settings tier); each is a file or dir loaded non-recursively. + ConfigPrompts []string + // CliPrompts holds --prompt-template paths (CLI tier, repeatable). + CliPrompts []string + // NoPromptTemplates disables all prompt-template discovery (global, project, + // settings, CLI); built-in slash commands are unaffected. Independent of + // --no-skills. + NoPromptTemplates bool +} diff --git a/pigo/internal/cli/tui/remotecontrol.go b/pigo/internal/cli/tui/remotecontrol.go new file mode 100644 index 0000000..eb0558e --- /dev/null +++ b/pigo/internal/cli/tui/remotecontrol.go @@ -0,0 +1,206 @@ +// This file wires the remote-control bridge (internal/remotecontrol, #442) into +// the full-screen TUI, the counterpart to internal/cli/repl/remotecontrol.go. +// It adds the "/remote-control" command that starts/stops an in-process +// HTTP+WebSocket server mirroring the session to a paired browser on the LAN, +// mirrors transcript output to that browser, surfaces browser-submitted prompts +// as a tea.Msg, and routes side-effect tool-call confirmations to the browser +// while a client is connected. +// +// The non-remote path is unchanged: when no session is active the mirror is a +// no-op, waitRemoteInput returns a nil Cmd, and buildConfig installs no +// BeforeToolCall (tools run under the up-front trust the TUI already grants). +package tui + +import ( + "context" + "fmt" + "strings" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/remotecontrol" + "github.com/smallnest/pigo/internal/trust" +) + +// remoteSession owns the running server + bridge for one /remote-control +// activation. It is stored on runSession (so buildConfig can reach it to install +// the confirm seam) and is nil until the command starts a session. +type remoteSession struct { + server *remotecontrol.Server + bridge *remotecontrol.Bridge + url string +} + +// hasClient reports whether a browser is currently paired and connected. +func (rs *remoteSession) hasClient() bool { + return rs != nil && rs.bridge != nil && rs.bridge.Enabled() +} + +// sendOutput mirrors session text to the remote browser. It records into the +// server's replay ring even when no client is connected, so a browser that pairs +// mid-session is replayed the recent scrollback. +func (rs *remoteSession) sendOutput(text string) { + if rs == nil || rs.server == nil || text == "" { + return + } + rs.server.SendOutput(text) +} + +// startRemote builds and starts the server+bridge, storing the session on the +// runSession. It returns the pairing URL, or an error if a server is already +// running or the listener could not bind. +func (s *runSession) startRemote() (string, error) { + if s.remote != nil { + return s.remote.url, fmt.Errorf("already running") + } + // Break the server↔bridge construction cycle: build the server (Sink), then + // the bridge over it, then route client frames back to the bridge. + srv := remotecontrol.NewServer(remotecontrol.Config{}, nil) + bridge := remotecontrol.NewBridge(srv) + srv.SetHandler(bridge) + url, err := srv.Start() + if err != nil { + return "", err + } + s.remote = &remoteSession{server: srv, bridge: bridge, url: url} + return url, nil +} + +// stopRemote shuts down the running server and clears the session. It is a no-op +// when remote control is off. +func (s *runSession) stopRemote() { + if s.remote == nil { + return + } + _ = s.remote.server.Stop(context.Background()) + s.remote = nil +} + +// remoteConfirmSeam builds the BeforeToolCall seam that routes side-effect +// tool-call confirmations to the paired browser while one is connected. When no +// browser is connected (or the tool is not side-effecting, or the cwd is +// trusted) it returns nil so the tool runs under the up-front trust the TUI +// grants — the non-remote behavior is unchanged. +// +// A ctx cancellation (interrupt) makes Confirm return remote=false, which is +// treated as a denial so an interrupted run does not silently proceed. +func remoteConfirmSeam(rs *remoteSession, mgr *trust.Manager, cwd string) agentcore.BeforeToolCallFunc { + return func(ctx context.Context, call agentcore.AgentToolCall) *agentcore.BeforeToolCallDecision { + if !rs.hasClient() || mgr == nil { + return nil + } + if !trust.SideEffectTools[call.Name] { + return nil + } + if mgr.IsTrusted(cwd) { + return nil + } + summary := trust.ToolCallSummary(call) + d, remote := rs.bridge.Confirm(ctx, call.Name, summary) + if !remote { + return blockRemoteToolCall(call, cwd) + } + if d.Always { + mgr.SetSessionTrust(cwd) + } + if !d.Approve { + return blockRemoteToolCall(call, cwd) + } + return nil + } +} + +func blockRemoteToolCall(call agentcore.AgentToolCall, cwd string) *agentcore.BeforeToolCallDecision { + msg := fmt.Sprintf("tool %q blocked: %s is not trusted (use /trust to trust this project)", call.Name, cwd) + return &agentcore.BeforeToolCallDecision{ + Block: true, + Content: &agentcore.ContentList{agentcore.NewTextContent(msg)}, + } +} + +// runRemoteControl handles the /remote-control command and its stop/status +// subcommands. It mutates m.session.remote, folds a system block into the +// transcript, and returns the listener Cmd (waitRemoteInput) on a successful +// start so browser-submitted prompts begin arriving. +func (m Model) runRemoteControl(line string) (tea.Model, tea.Cmd) { + m.transcript.addUser(line) + m.input.Clear() + m.menu.close() + defer m.relayout() + + if m.session == nil { + m.transcript.addSystem("(remote control unavailable: no active session)") + return m, nil + } + arg := strings.TrimSpace(strings.TrimPrefix(line, "/remote-control")) + switch arg { + case "stop": + if m.session.remote == nil { + m.transcript.addSystem("remote control is not running") + return m, nil + } + m.session.stopRemote() + m.transcript.addSystem("remote control stopped") + return m, nil + case "status": + if m.session.remote == nil { + m.transcript.addSystem("remote control: off") + return m, nil + } + state := "waiting for a browser to connect" + if m.session.remote.hasClient() { + state = "browser connected" + } + m.transcript.addSystem(fmt.Sprintf("remote control: on (%s)\n %s", state, m.session.remote.url)) + return m, nil + case "": + if m.session.remote != nil { + m.transcript.addSystem("remote control already running:\n " + m.session.remote.url) + return m, nil + } + url, err := m.session.startRemote() + if err != nil { + m.transcript.addSystem("remote control: " + err.Error()) + return m, nil + } + var b strings.Builder + fmt.Fprintf(&b, "Remote control started. Open this URL on a device on the same network:\n\n %s\n", url) + if qr, qerr := remotecontrol.Render(url); qerr == nil { + b.WriteString("\n" + qr) + } + b.WriteString("\nRun /remote-control stop to end the session.") + m.transcript.addSystem(b.String()) + return m, m.waitRemoteInput() + default: + m.transcript.addSystem("usage: /remote-control [stop|status]") + return m, nil + } +} + +// remoteEcho mirrors visible transcript text to the paired browser. It is a +// no-op when remote control is off, so callers can invoke it unconditionally at +// each point the transcript gains content. +func (m Model) remoteEcho(text string) { + if m.session != nil && m.session.remote != nil { + m.session.remote.sendOutput(text) + } +} + +// waitRemoteInput returns a tea.Cmd that blocks on the bridge's remote-input +// channel and emits one remoteInputMsg per browser submission. The Update loop +// re-issues it after each so successive prompts keep arriving. It returns nil +// when remote control is off, which stops the listener. +func (m Model) waitRemoteInput() tea.Cmd { + if m.session == nil || m.session.remote == nil || m.session.remote.bridge == nil { + return nil + } + ch := m.session.remote.bridge.RemoteInput() + return func() tea.Msg { + text, ok := <-ch + if !ok { + return nil + } + return remoteInputMsg{text: text} + } +} diff --git a/pigo/internal/cli/tui/remotecontrol_test.go b/pigo/internal/cli/tui/remotecontrol_test.go new file mode 100644 index 0000000..cb53269 --- /dev/null +++ b/pigo/internal/cli/tui/remotecontrol_test.go @@ -0,0 +1,145 @@ +package tui + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// newRemoteTestSession builds a fresh run session over a temp-dir store for the +// remote-control lifecycle tests (no resume, no tools). +func newRemoteTestSession(t *testing.T) *runSession { + t.Helper() + store := newTestStore(t) + s, _, err := newRunSessionWithStore(store, Options{}) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + return s +} + +// TestStartStopRemote covers the start→already-running→stop lifecycle on +// runSession: startRemote binds a listener and stores the session, a second +// start reports "already running" without replacing it, and stopRemote clears +// the session so a later start can rebind. +func TestStartStopRemote(t *testing.T) { + s := newRemoteTestSession(t) + if s.remote != nil { + t.Fatal("remote should be nil before start") + } + + url, err := s.startRemote() + if err != nil { + t.Fatalf("startRemote: %v", err) + } + if url == "" { + t.Fatal("startRemote returned empty url") + } + if s.remote == nil { + t.Fatal("remote should be set after start") + } + first := s.remote + + // A second start is a no-op that reports the existing url and an error, + // leaving the running session untouched. + url2, err := s.startRemote() + if err == nil { + t.Error("second startRemote should report already-running error") + } + if url2 != url { + t.Errorf("second startRemote url = %q, want %q", url2, url) + } + if s.remote != first { + t.Error("second startRemote must not replace the running session") + } + + s.stopRemote() + if s.remote != nil { + t.Error("remote should be nil after stop") + } + + // Stop again is a no-op. + s.stopRemote() + + // After stopping, a fresh start rebinds cleanly. + if _, err := s.startRemote(); err != nil { + t.Fatalf("restart after stop: %v", err) + } + s.stopRemote() +} + +// TestBuildConfigInstallsRemoteSeam asserts buildConfig only installs the +// BeforeToolCall confirm seam while a remote session is present: off by default +// (up-front trust unchanged), wired once /remote-control is running. +func TestBuildConfigInstallsRemoteSeam(t *testing.T) { + s := newRemoteTestSession(t) + + if cfg := s.buildConfig(); cfg.Batch.ToolExecutorConfig.BeforeToolCall != nil { + t.Error("BeforeToolCall should be nil when remote control is off") + } + + if _, err := s.startRemote(); err != nil { + t.Fatalf("startRemote: %v", err) + } + defer s.stopRemote() + + if cfg := s.buildConfig(); cfg.Batch.ToolExecutorConfig.BeforeToolCall == nil { + t.Error("BeforeToolCall should be installed when remote control is on") + } +} + +// TestRemoteConfirmSeamAllowsWhenNoClient verifies the confirm seam is a no-op +// (returns nil = allow under up-front trust) when no browser is connected, so a +// running-but-unpaired server never blocks tool calls. +func TestRemoteConfirmSeamAllowsWhenNoClient(t *testing.T) { + s := newRemoteTestSession(t) + if _, err := s.startRemote(); err != nil { + t.Fatalf("startRemote: %v", err) + } + defer s.stopRemote() + + // No client is paired, so hasClient() is false and the seam must allow. + seam := remoteConfirmSeam(s.remote, nil, "/tmp/project") + if d := seam(t.Context(), agentcore.AgentToolCall{Name: "bash"}); d != nil { + t.Errorf("seam should allow (nil) with no client, got %+v", d) + } +} + +// TestRemoteInputIgnoredWhileRunning routes a remoteInputMsg into an idle vs a +// running Model: while a run is in flight the prompt is not started (a busy note +// is shown instead), and the listener is always re-issued so later submissions +// keep arriving. +func TestRemoteInputIgnoredWhileRunning(t *testing.T) { + s := newRemoteTestSession(t) + if _, err := s.startRemote(); err != nil { + t.Fatalf("startRemote: %v", err) + } + defer s.stopRemote() + + m := NewModel(Options{}) + m.session = s + m.running = true + + updated, _ := m.Update(remoteInputMsg{text: "do something"}) + got := updated.(Model) + if !hasSystemBlockContaining(got.transcript, "a run is in progress") { + t.Errorf("expected a busy note in the transcript while running, blocks=%v", + blockTexts(got.transcript)) + } + // A run in progress must not consume the remote prompt as a new turn. + if !got.running { + t.Error("model should still be running; the remote prompt must not start a turn") + } +} + +// hasSystemBlockContaining reports whether any transcript block's text contains +// sub (case-insensitive substring over the rendered block texts). +func hasSystemBlockContaining(t transcript, sub string) bool { + for _, s := range blockTexts(t) { + if strings.Contains(strings.ToLower(s), strings.ToLower(sub)) { + return true + } + } + return false +} diff --git a/pigo/internal/cli/tui/run.go b/pigo/internal/cli/tui/run.go new file mode 100644 index 0000000..ca11da3 --- /dev/null +++ b/pigo/internal/cli/tui/run.go @@ -0,0 +1,24 @@ +package tui + +import ( + tea "charm.land/bubbletea/v2" +) + +// Run starts the full-screen TUI and blocks until the user quits (Ctrl+C / +// Ctrl+D) or the program errors. It is the alt-screen counterpart to repl.Run: +// cmd/pigo's dispatch calls it on the (no prompt + TTY + no --no-tui) path and +// maps its error to the process exit code. The alt-screen is entered/left via +// the View returned by the root Model, so a clean return here restores the +// terminal to the user's prior scrollback. +func Run(opts Options) error { + // Assemble the session (store, resume-or-fresh context, live config) before + // entering the alt-screen, mirroring repl.Run: a store/resume failure is a + // clean pre-launch error rather than a broken interactive session. + s, history, err := newRunSession(opts) + if err != nil { + return err + } + p := tea.NewProgram(NewModel(opts).withSession(s, history)) + _, err = p.Run() + return err +} diff --git a/pigo/internal/cli/tui/selection.go b/pigo/internal/cli/tui/selection.go new file mode 100644 index 0000000..03dd124 --- /dev/null +++ b/pigo/internal/cli/tui/selection.go @@ -0,0 +1,112 @@ +package tui + +import ( + "strings" + + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/ansi" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// This file implements mouse text selection over the rendered shell. Because the +// model paints the whole screen as one string (transcript + menu + input + +// status), a selection is expressed in screen cells (0-based, top-left origin) +// and spans any region uniformly — dragging across transcript output or the +// input line both work. The left mouse button starts a selection at the press +// cell and extends it on drag; the selection persists after release so Ctrl+C +// can copy it. A plain click (no drag) leaves an empty selection, which clears +// any prior highlight and lets Ctrl+C fall back to its interrupt/quit role. + +// maxCol is a sentinel column that reaches past the end of any rendered row, so +// a multi-row selection's interior rows select through to their line end. +const maxCol = 1 << 30 + +// point is a screen cell: x is the column, y the row, both 0-based from the +// top-left of the rendered shell. +type point struct{ x, y int } + +// selection is an in-progress or completed text selection. anchor is where the +// drag began and cursor is the latest drag point; active is set between the +// initial press and the next fresh press. It is a plain value so the Model can +// hold and copy it cheaply. +type selection struct { + active bool + anchor point + cursor point +} + +// empty reports whether the selection covers no cells — either inactive or a +// bare click where the cursor never moved off the anchor. Ctrl+C treats an empty +// selection as "nothing to copy" and keeps its interrupt/quit behavior. +func (s selection) empty() bool { + return !s.active || s.anchor == s.cursor +} + +// ordered returns the selection endpoints in reading order (top-to-bottom, and +// left-to-right within a row), so callers can walk rows start.y..end.y without +// re-checking which of anchor/cursor came first. +func (s selection) ordered() (start, end point) { + a, c := s.anchor, s.cursor + if a.y > c.y || (a.y == c.y && a.x > c.x) { + return c, a + } + return a, c +} + +// rowRange computes the selected column span [c0, c1) on screen row y for a +// selection running start..end. Interior rows of a multi-row selection run from +// column 0 through the line end (maxCol); the first row starts at start.x and +// the last ends at end.x. ok is false when y falls outside the selection. +func rowRange(start, end point, y int) (c0, c1 int, ok bool) { + if y < start.y || y > end.y { + return 0, 0, false + } + c0, c1 = 0, maxCol + if y == start.y { + c0 = start.x + } + if y == end.y { + c1 = end.x + } + if c0 < 0 { + c0 = 0 + } + if c1 < c0 { + c1 = c0 + } + return c0, c1, true +} + +// selectRow walks one rendered row (ANSI stripped to plain cells) and returns +// both the row with the selected span visually highlighted and the selected +// text itself. Columns are measured in display cells (ui.Width) so double-width +// runes are never split. The highlight is applied over plain text — the row's +// original coloring is dropped on the intersected row while a selection is live, +// which keeps the overlay ANSI-safe without parsing the embedded escapes. +func selectRow(row string, c0, c1 int, hi lipgloss.Style) (highlighted, text string) { + plain := ansi.Strip(row) + + var out, sel, run strings.Builder + flush := func() { + if run.Len() > 0 { + out.WriteString(hi.Render(run.String())) + run.Reset() + } + } + + col := 0 + for _, r := range plain { + w := ui.Width(string(r)) + if col >= c0 && col < c1 { + run.WriteRune(r) + sel.WriteRune(r) + } else { + flush() + out.WriteRune(r) + } + col += w + } + flush() + return out.String(), sel.String() +} diff --git a/pigo/internal/cli/tui/selection_test.go b/pigo/internal/cli/tui/selection_test.go new file mode 100644 index 0000000..1423244 --- /dev/null +++ b/pigo/internal/cli/tui/selection_test.go @@ -0,0 +1,59 @@ +package tui + +import ( + "testing" + + "charm.land/lipgloss/v2" +) + +// TestRowRange checks the per-row column span for single- and multi-row +// selections, and that rows outside the range report ok=false. +func TestRowRange(t *testing.T) { + start, end := point{3, 1}, point{7, 3} + + if _, _, ok := rowRange(start, end, 0); ok { + t.Error("row above the selection should not be selected") + } + if c0, c1, ok := rowRange(start, end, 1); !ok || c0 != 3 || c1 != maxCol { + t.Errorf("first row = (%d,%d,%v), want (3,maxCol,true)", c0, c1, ok) + } + if c0, c1, ok := rowRange(start, end, 2); !ok || c0 != 0 || c1 != maxCol { + t.Errorf("interior row = (%d,%d,%v), want (0,maxCol,true)", c0, c1, ok) + } + if c0, c1, ok := rowRange(start, end, 3); !ok || c0 != 0 || c1 != 7 { + t.Errorf("last row = (%d,%d,%v), want (0,7,true)", c0, c1, ok) + } + if _, _, ok := rowRange(start, end, 4); ok { + t.Error("row below the selection should not be selected") + } + + // A single-row selection uses [start.x, end.x). + if c0, c1, ok := rowRange(point{2, 5}, point{9, 5}, 5); !ok || c0 != 2 || c1 != 9 { + t.Errorf("single row = (%d,%d,%v), want (2,9,true)", c0, c1, ok) + } +} + +// TestSelectRowExtracts verifies the selected text is the column-clipped slice +// of the row, measured in display cells so CJK is never split, and that ANSI in +// the source row is stripped before slicing. +func TestSelectRowExtracts(t *testing.T) { + hi := lipgloss.NewStyle().Reverse(true) + + if _, text := selectRow("hello world", 0, 5, hi); text != "hello" { + t.Errorf("selected %q, want %q", text, "hello") + } + if _, text := selectRow("hello world", 6, maxCol, hi); text != "world" { + t.Errorf("selected %q, want %q", text, "world") + } + + // ANSI coloring in the source is stripped before the selection is measured. + styled := lipgloss.NewStyle().Foreground(lipgloss.Color("42")).Render("hello") + if _, text := selectRow(styled, 0, maxCol, hi); text != "hello" { + t.Errorf("selected %q from styled row, want %q", text, "hello") + } + + // CJK counts as two columns: selecting the first two cells yields one rune. + if _, text := selectRow("你好ab", 0, 2, hi); text != "你" { + t.Errorf("selected %q, want %q (double-width clipped on a cell boundary)", text, "你") + } +} diff --git a/pigo/internal/cli/tui/session.go b/pigo/internal/cli/tui/session.go new file mode 100644 index 0000000..f085d7a --- /dev/null +++ b/pigo/internal/cli/tui/session.go @@ -0,0 +1,491 @@ +// This file binds the full-screen TUI to the real agent run seam and the local +// session store (US-009, FR-16/17). It is the TUI counterpart to the REPL's +// replDeps + streamRun + cli.PersistTurn plumbing (internal/cli/repl): it +// assembles an AgentContext + RunConfig from the model's Options, feeds them to +// the event bridge (bridge.go's startRun → runtime.StartRun/DrainStream), and +// persists the growing conversation to ~/.pigo/sessions after each turn. +// +// It deliberately imports the SHARED lower-level packages the REPL also uses +// (session, runtime, provider, cli, cli/run, cli/headless, cli/ui) rather than +// the repl package itself, so the two entry paths share one store and one +// run-config shape without an import cycle (repl and tui are siblings; prompts +// imports tui, so tui must not reach back into repl/prompts). +package tui + +import ( + "context" + "fmt" + "os" + "time" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/headless" + "github.com/smallnest/pigo/internal/cli/run" + "github.com/smallnest/pigo/internal/cli/ui" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/hooks" + "github.com/smallnest/pigo/internal/memory" + "github.com/smallnest/pigo/internal/plugin" + "github.com/smallnest/pigo/internal/provider" + "github.com/smallnest/pigo/internal/runtime" + "github.com/smallnest/pigo/internal/session" + "github.com/smallnest/pigo/internal/trust" +) + +// runSession holds the assembled per-session state for a TUI run: the persisted +// store + header, the growing conversation context, the live (mutable) run +// config, and the tool/credential collaborators. It mirrors the subset of +// repl.replDeps the TUI needs, and owns the same session-tree cursor bookkeeping +// (curLeaf / persisted) so each turn is persisted as a branch rather than a +// flattening rewrite. +type runSession struct { + store *session.Store + header session.SessionHeader + agentCtx *agentcore.AgentContext + live *cli.LiveConfig + reg *agenttool.ToolRegistry + reminders *runtime.ReminderRegistry + creds *provider.CredentialStore + + // cwd is the directory pigo was launched in, captured once at session + // assembly. It is the trust key and the /status environment display. + cwd string + // trust persists project-trust decisions (US-018, #134). It is nil when + // trust is disabled (store could not be loaded / no cwd); when nil /status + // reports "disabled" and the trust-gated hook layer is skipped. + trust *trust.Manager + // slash is the shared slash-command registry the TUI consults exactly as the + // REPL does. It is assembled per-session against the live config (withSession + // rebinds the model's registry to this one) so /model switches reach it. + slash *runtime.SlashRegistry + // telemetry holds the retained per-run telemetry events (US-001, #291) and + // the cumulative accumulator that sums metrics across all runs in the + // session. The run loop folds each run's TelemetryEvent into it; /status + // reads it back through the Host contract. + telemetry *cli.TelemetryHolder + + // memoryRoot is the persistent-memory Store root (empty when memory is + // disabled). It routes auto-compaction checkpoints and /rebuild recovery to + // /sessions//, the canonical checkpoint location. + memoryRoot string + // memstore is the live persistent-memory Store (nil when memory is disabled). + // It lets /memory inspect entry counts without re-opening the database. + memstore *memory.Store + + // dispatcher is the session's hook dispatcher, nil when no hooks are + // configured (FR-18). hookDeps carries the session id / project dir stamped + // onto every HookInput and hook process environment. + dispatcher *hooks.Dispatcher + hookDeps run.HookDeps + // onEvent is the observer chain delivered to every run: the plugin notifier + // (US-017) with the SessionEnd/PreCompact hook notifier chained after it. + onEvent func(agentcore.AgentEvent) + + // curLeaf is the id of the on-disk entry the next turn descends from; persisted + // is the number of agentCtx.Messages already written. persist() appends only + // Messages[persisted:] as a branch from curLeaf (see cli.PersistTurn). + curLeaf string + persisted int + + // compacted is set when the run loop compacted the context (CompactionEvent): + // compaction rewrites Messages into a summary + recent tail, which both shrinks + // the slice below persisted (so an incremental Messages[persisted:] would panic) + // and invalidates the branch prefix. persist() honors this by re-saving the + // flattened context linearly and resetting the branch cursor, then clears it. + compacted bool + + // cancelRun cancels the in-flight run's context; startRun sets it and the + // two-stage interrupt (Model.interruptFn → interrupt) calls it. It is nil + // before the first run and after a run is cancelled. + cancelRun context.CancelFunc + + // lastBtw is the /btw side thread's context from this process and + // lastBtwBase the background-message index it diverged from. Both are + // carried on the Host contract for parity with the REPL; the TUI does not + // run /btw today, so they stay nil/0. + lastBtw *agentcore.AgentContext + lastBtwBase int + + // remote owns the running remote-control server+bridge (remotecontrol.go), + // nil when /remote-control is off. buildConfig reads it to install the remote + // confirm seam so risky tool calls route to the paired browser while connected. + remote *remoteSession +} + +// newRunSession assembles the run session from the resolved Options, opening the +// shared ~/.pigo/sessions store. When Options carries a ResumeID it loads that +// session's entries and rebuilds the context (the returned history seeds the +// replayed transcript); otherwise it starts a fresh session with a new header. +// It is the production entry; newRunSessionWithStore holds the store-agnostic +// core so tests can drive it against a temp-dir store. +func newRunSession(opts Options) (*runSession, []agentcore.Message, error) { + store, err := headless.SessionStore() + if err != nil { + return nil, nil, err + } + return newRunSessionWithStore(store, opts) +} + +// newRunSessionWithStore is the store-agnostic core of newRunSession: given an +// already-opened store it resolves resume-vs-fresh, builds the live config and +// collaborators, and returns the session plus the resumed history (nil for a +// fresh session). +func newRunSessionWithStore(store *session.Store, opts Options) (*runSession, []agentcore.Message, error) { + creds := provider.NewCredentialStore(nil) + creds.SetOverride(opts.ProviderName, opts.APIKey) + + // cwd is the launch directory, captured once: it stamps fresh sessions, is + // the trust key, and feeds /status's environment section and the hook layer. + cwd, _ := os.Getwd() + + now := time.Now().UTC() + var ( + agentCtx *agentcore.AgentContext + header session.SessionHeader + history []agentcore.Message + curLeaf string + ) + if opts.ResumeID != "" { + h, entries, err := store.LoadEntries(opts.ResumeID) + if err != nil { + return nil, nil, err + } + msgs := make(agentcore.MessageList, len(entries)) + for i, e := range entries { + msgs[i] = e.Message + } + if len(entries) > 0 { + curLeaf = entries[len(entries)-1].ID + } + header = h + sysPrompt := h.SystemPrompt + if sysPrompt == "" { + sysPrompt = opts.SysPrompt + } + agentCtx = &agentcore.AgentContext{SystemPrompt: sysPrompt, Messages: msgs, Tools: opts.Tools} + history = msgs + } else { + agentCtx = &agentcore.AgentContext{SystemPrompt: opts.SysPrompt, Tools: opts.Tools} + // Stamp the launch directory onto a fresh session (#526/#524) so the + // session is attributed to a project and a later /dream pass can distill it + // under the right scope, mirroring headless/REPL. An unresolvable cwd + // yields "" (session stays unattributed) rather than aborting. + header = session.SessionHeader{ + ID: session.NewID(now), + CreatedAt: now, + UpdatedAt: now, + Model: opts.Model, + Provider: opts.ProviderName, + SystemPrompt: opts.SysPrompt, + Cwd: cwd, + } + } + + live := &cli.LiveConfig{ + Model: opts.Model, + ProviderName: opts.ProviderName, + Provider: opts.Provider, + BaseURL: opts.BaseURL, + Protocol: opts.Protocol, + ThinkingLevel: opts.ThinkingLevel, + ContextWindow: cli.DefaultContextWindow, + } + + // Project trust (US-018, #134): load the persisted trust store for the + // launch directory, mirroring the REPL. A load failure (or an unresolvable + // cwd) is non-fatal: trust is disabled (mgr stays nil) and the TUI still + // runs — the store is surfaced rather than silently overwritten. + mgr, mgrErr := trust.NewManager(trust.DefaultPath()) + if mgrErr != nil { + fmt.Fprintf(os.Stderr, "pigo: trust store unavailable, trust disabled: %v\n", mgrErr) + mgr = nil + } + if cwd == "" && mgr != nil { + fmt.Fprintf(os.Stderr, "pigo: cannot resolve working directory, trust disabled\n") + mgr = nil + } + + s := &runSession{ + store: store, + header: header, + agentCtx: agentCtx, + live: live, + reg: run.ToolRegistry(opts.Tools), + reminders: run.TodoReminders(opts.Tools), + creds: creds, + cwd: cwd, + trust: mgr, + slash: newSlashRegistry(opts, live), + telemetry: cli.NewTelemetryHolder(), + curLeaf: curLeaf, + persisted: len(history), + memoryRoot: run.MemoryRootFromTools(opts.Tools), + memstore: run.MemoryStoreFromTools(opts.Tools), + } + // /trust is a per-session command (its closure captures mgr + cwd), so it is + // registered here rather than in newSlashRegistry. A nil mgr is a no-op. + trust.RegisterCommand(s.slash, mgr, cwd) + + // Wire hooks uniformly with every other driver (#425): resolve the trust-gated + // hook set, build the dispatcher, dispatch SessionStart once, and compose the + // SessionEnd/PreCompact observer with the plugin notifier. Trust is granted by + // --approve (Options.Approve) or the shared trust store; project-layer hooks + // only apply when trusted (FR-14). A malformed hook layer disables hooks with a + // warning rather than failing the TUI launch. + s.hookDeps = run.HookDeps{SessionID: header.ID, ProjectDir: cwd, WarnLog: os.Stderr} + trusted := opts.Approve || (mgr != nil && mgr.IsTrusted(cwd)) + var baseOnEvent func(agentcore.AgentEvent) + if n := plugin.NewEventNotifier(opts.Plugins, os.Stderr); n != nil { + baseOnEvent = n.Handle + } + if set, err := run.ResolveHookSet(cwd, trusted); err != nil { + fmt.Fprintf(os.Stderr, "pigo: hooks disabled: %v\n", err) + s.onEvent = baseOnEvent + } else if d := run.BuildDispatcher(set, s.hookDeps); d != nil { + s.dispatcher = d + if s.reminders == nil { + s.reminders = runtime.NewReminderRegistry() + } + ssCfg := runtime.RunConfig{Reminders: s.reminders} + run.DispatchSessionStart(context.Background(), d, &ssCfg, s.hookDeps, sessionStartSource(opts)) + s.reminders = ssCfg.Reminders + n := hooks.NewHookNotifier(d, s.hookDeps.SessionID, s.hookDeps.ProjectDir) + s.onEvent = chainTUIEvent(baseOnEvent, n.Handle) + } else { + s.onEvent = baseOnEvent + } + return s, history, nil +} + +// sessionStartSource maps the resolved run options to the SessionStart source +// tag: "resume" when continuing an existing session, "startup" otherwise. +func sessionStartSource(opts Options) string { + if opts.ResumeID != "" { + return "resume" + } + return "startup" +} + +// chainTUIEvent composes the plugin notifier with the hook notifier into one +// observer; a nil operand is identity. +func chainTUIEvent(prev, next func(agentcore.AgentEvent)) func(agentcore.AgentEvent) { + if prev == nil { + return next + } + if next == nil { + return prev + } + return func(ev agentcore.AgentEvent) { + prev(ev) + next(ev) + } +} + +// buildConfig assembles the RunConfig for one turn from the live config and +// collaborators. It replicates repl.streamRun's assembly (same LoopConfig fields, +// tool registry and reminders) minus the interactive trust confirmation hook: the +// TUI has no stdin prompt to confirm side-effect tool calls on, so tools run +// under the trust granted up front by --approve (Options.Approve) rather than a +// per-call BeforeToolCall prompt. The stream fn is derived from the live provider +// and the API key resolved through the credential store, exactly as the REPL does. +func (s *runSession) buildConfig() runtime.RunConfig { + cfg := runtime.RunConfig{ + LoopConfig: runtime.LoopConfig{ + Model: s.live.Model, + Provider: s.live.ProviderName, + ThinkingLevel: s.live.ThinkingLevel, + Stream: provider.StreamFnFromProvider(s.live.Provider), + GetAPIKey: s.creds.GetAPIKey, + ContextWindow: s.live.ContextWindow, + Compaction: compaction.DefaultCompactionSettings, + }, + Batch: agenttool.BatchConfig{ + ToolExecutorConfig: agenttool.ToolExecutorConfig{ + Registry: s.reg, + }, + }, + Reminders: s.reminders, + SessionID: s.header.ID, + MemoryRoot: s.memoryRoot, + } + // Per-turn wiring of the tool-execution + Stop seams; nil dispatcher is a + // no-op so the hot path pays nothing when no hooks are configured (FR-18). + if s.dispatcher != nil { + run.InstallSeams(&cfg, s.dispatcher, s.hookDeps) + } + // When remote control is active, route side-effect tool-call confirmations to + // the paired browser (no-op when no client is connected or the cwd is trusted, + // so the non-remote path is unchanged). The trust manager is read from the + // shared store; a nil manager disables the seam. + if s.remote != nil { + if mgr, err := trust.NewManager(trust.DefaultPath()); err == nil { + cfg.Batch.ToolExecutorConfig.BeforeToolCall = remoteConfirmSeam(s.remote, mgr, s.hookDeps.ProjectDir) + } + } + return cfg +} + +// rebuildDoneMsg reports the outcome of a manual /rebuild to the model: summary +// is the status line to show in the transcript, err is set when the rebuild +// failed (the context is then left unchanged). +type rebuildDoneMsg struct { + summary string + err error +} + +// rebuildCmd runs a context rebuild off the tea loop (the no-checkpoint fallback +// makes a summarization LLM call, so it must not block the UI goroutine) and +// yields a rebuildDoneMsg the model folds into the transcript. It mirrors the +// REPL's runManualRebuild. +func (s *runSession) rebuildCmd() tea.Cmd { + return func() tea.Msg { + summary, err := s.rebuild() + return rebuildDoneMsg{summary: summary, err: err} + } +} + +// rebuild reconstructs the shared context from the session's persisted checkpoint +// (collapsing the pre-watermark prefix to the checkpoint summary and preserving +// the recent tail verbatim), falling back to lossy compaction when no checkpoint +// exists. It replaces agentCtx.Messages in place on success and flags compacted +// so persist() re-saves the flattened context linearly (as after a /compact). +func (s *runSession) rebuild() (string, error) { + msgs := s.agentCtx.Messages + before := compaction.EstimateContextTokens(msgs).Tokens + // Checkpoints live under /sessions//; recover from the same + // root the loop writes to. Empty when memory is disabled — RebuildFromCheckpoint + // then falls back to lossy compaction. + memoryRoot := s.memoryRoot + cfg := s.buildConfig() + res, err := runtime.RebuildFromCheckpoint(context.Background(), msgs, s.header.ID, memoryRoot, &cfg, nil) + if err != nil { + return "", err + } + if res.NoOp { + return fmt.Sprintf("nothing to rebuild (%d tokens, %d messages)", before, len(msgs)), nil + } + s.agentCtx.Messages = res.Messages + s.compacted = true + source := "checkpoint" + if !res.FromCheckpoint { + source = "compaction (no checkpoint)" + } + return fmt.Sprintf("context rebuilt from %s: %d → %d tokens, collapsed %d messages, kept %d", + source, res.TokensBefore, res.TokensAfter, res.SummarizedCount, res.KeptCount), nil +} + + +// prompt to the growing context as a user message, then hands the context and a +// freshly-built config to the event bridge (bridge.startRun → runtime.StartRun + +// DrainStream on a goroutine), returning the bridge channel and the first +// waitForEvent Cmd so Update can pump the run's events. The context grows in +// place (agentCtx is a pointer), so the next turn continues the conversation. +func (s *runSession) startRun(prompt string) (chan tea.Msg, tea.Cmd) { + content, err := ui.BuildUserContent(prompt) + if err != nil { + // A malformed image reference must not swallow the turn: fall back to the + // raw prompt as plain text so the run still starts. + content = agentcore.ContentList{agentcore.NewTextContent(prompt)} + } + // UserPromptSubmit runs before the prompt is committed to the context: a block + // aborts the turn (emitting a runEndMsg carrying the reason) without leaving a + // dangling user message; additionalContext is injected into this turn only. + if s.dispatcher != nil { + pc := runtime.RunConfig{Reminders: s.reminders} + if block, reason := run.DispatchUserPromptSubmit(context.Background(), s.dispatcher, &pc, s.hookDeps, prompt); block { + ch := newEventChan() + go func() { ch <- runEndMsg{err: fmt.Errorf("prompt blocked by hook: %s", reason)} }() + return ch, waitForEvent(ch) + } + s.reminders = pc.Reminders + } + s.agentCtx.Messages = append(s.agentCtx.Messages, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: content, + }) + // Use a cancellable context so the two-stage interrupt (FR-14) can stop this + // run: cancelling propagates through StartRun/DrainStream, which then emits a + // runEndMsg and the model returns to idle. + ctx, cancel := context.WithCancel(context.Background()) + s.cancelRun = cancel + return startRun(ctx, s.agentCtx, s.buildConfig(), s.onEvent) +} + +// interrupt cancels the in-flight run, if any. It is bound to Model.interruptFn +// by withSession so pressing Esc / Ctrl+C while running stops the current run +// instead of quitting the program (FR-14). Safe to call when no run is active. +func (s *runSession) interrupt() { + if s.cancelRun != nil { + s.cancelRun() + } +} + +// persist writes the messages produced since the last persist as a new branch +// descending from the active leaf, advancing the leaf and the persisted cursor. +// It mirrors cli.PersistTurn: growing the on-disk tree with AppendBranch (rather +// than a linear rewrite) keeps history intact. A no-op when nothing new was +// produced, so an idle turn-end never regenerates entry ids. +func (s *runSession) persist() error { + // A compaction during the run rewrote Messages into a summary + recent tail, + // so the append-a-tail branch model no longer holds: the prefix changed and + // the slice may be shorter than persisted. Re-save the flattened context + // linearly and reset the branch cursor to the new leaf, mirroring the REPL's + // /compact handling. + if s.compacted || s.persisted > len(s.agentCtx.Messages) { + s.header.UpdatedAt = time.Now().UTC() + s.header.Model = s.live.Model + s.header.Provider = s.live.ProviderName + if err := s.store.Save(s.header, s.agentCtx.Messages); err != nil { + return err + } + s.persisted = len(s.agentCtx.Messages) + s.curLeaf = "" + if _, entries, err := s.store.LoadEntries(s.header.ID); err == nil && len(entries) > 0 { + s.curLeaf = entries[len(entries)-1].ID + } + s.compacted = false + return nil + } + tail := s.agentCtx.Messages[s.persisted:] + if len(tail) == 0 { + return nil + } + s.header.UpdatedAt = time.Now().UTC() + s.header.Model = s.live.Model + s.header.Provider = s.live.ProviderName + leaf, err := s.store.AppendBranch(s.header, s.curLeaf, tail) + if err != nil { + return err + } + s.curLeaf = leaf + s.persisted = len(s.agentCtx.Messages) + return nil +} + +// seedTranscript replays a resumed session's prior messages into the transcript +// so the user sees the conversation so far before re-prompting (the TUI analogue +// of repl.replayTranscript). User and assistant text become their respective +// blocks; assistant tool calls render as system lines (tool cards land in #389). +// Tool-result messages are omitted here — their content is echoed live during a +// run, and replaying raw results would clutter the resumed view. +func seedTranscript(t *transcript, history []agentcore.Message) { + for _, m := range history { + switch msg := m.(type) { + case agentcore.UserMessage: + if text := agentcore.ContentToText(msg.Content); text != "" { + t.addUser(text) + } + case agentcore.AssistantMessage: + if text := agentcore.ContentToText(msg.Content); text != "" { + t.finalizeTurn(msg) + } + for _, c := range msg.ToolCalls() { + t.addSystem("· " + c.Name) + } + } + } +} diff --git a/pigo/internal/cli/tui/session_test.go b/pigo/internal/cli/tui/session_test.go new file mode 100644 index 0000000..6659851 --- /dev/null +++ b/pigo/internal/cli/tui/session_test.go @@ -0,0 +1,210 @@ +package tui + +import ( + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/session" +) + +// newTestStore opens a session store rooted at a temp dir so persistence/resume +// can be exercised without touching ~/.pigo. +func newTestStore(t *testing.T) *session.Store { + t.Helper() + store, err := session.NewStore(t.TempDir()) + if err != nil { + t.Fatalf("NewStore: %v", err) + } + return store +} + +// saveSession writes a linear session with the given messages and returns its id. +func saveSession(t *testing.T, store *session.Store, msgs agentcore.MessageList) string { + t.Helper() + now := time.Now().UTC() + header := session.SessionHeader{ + ID: session.NewID(now), + CreatedAt: now, + UpdatedAt: now, + Model: "test-model", + Provider: "test-provider", + } + if err := store.Save(header, msgs); err != nil { + t.Fatalf("Save: %v", err) + } + return header.ID +} + +// TestResumeSeedsTranscript constructs a session with a few messages, resumes it +// through newRunSessionWithStore, seeds a transcript with the returned history, +// and asserts the initial transcript blocks carry those messages (FR-16 resume). +func TestResumeSeedsTranscript(t *testing.T) { + store := newTestStore(t) + id := saveSession(t, store, agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hello, world")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("hello back")}}, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("second question")}}, + }) + + s, history, err := newRunSessionWithStore(store, Options{ResumeID: id}) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + if len(history) != 3 { + t.Fatalf("history len = %d, want 3", len(history)) + } + // The persisted cursor must cover the full resumed history so the first new + // turn appends only fresh messages, not a re-save of history. + if s.persisted != 3 { + t.Errorf("persisted = %d, want 3", s.persisted) + } + if s.curLeaf == "" { + t.Error("curLeaf should be the resumed leaf, got empty") + } + + tr := newTranscript(DefaultTheme()) + seedTranscript(&tr, history) + + wantTexts := []string{"hello, world", "hello back", "second question"} + if len(tr.blocks) != len(wantTexts) { + t.Fatalf("transcript blocks = %d, want %d", len(tr.blocks), len(wantTexts)) + } + for i, want := range wantTexts { + if tr.blocks[i].text != want { + t.Errorf("block[%d] = %q, want %q", i, tr.blocks[i].text, want) + } + } +} + +// TestBuildConfigAssembly asserts the run-config assembly maps the live config +// onto RunConfig without a live provider: the model/provider/thinking/window +// fields flow through, compaction is enabled, and the tool registry is wired. +func TestBuildConfigAssembly(t *testing.T) { + store := newTestStore(t) + s, _, err := newRunSessionWithStore(store, Options{ + Model: "opus-test", + ProviderName: "anthropic", + ThinkingLevel: agentcore.ThinkingLevel("high"), + }) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + + cfg := s.buildConfig() + if cfg.Model != "opus-test" { + t.Errorf("cfg.Model = %q, want opus-test", cfg.Model) + } + if cfg.Provider != "anthropic" { + t.Errorf("cfg.Provider = %q, want anthropic", cfg.Provider) + } + if cfg.ThinkingLevel != agentcore.ThinkingLevel("high") { + t.Errorf("cfg.ThinkingLevel = %q, want high", cfg.ThinkingLevel) + } + if cfg.ContextWindow <= 0 { + t.Errorf("cfg.ContextWindow = %d, want a positive default", cfg.ContextWindow) + } + if !cfg.Compaction.Enabled { + t.Error("cfg.Compaction.Enabled = false, want true (DefaultCompactionSettings)") + } + if cfg.Batch.Registry == nil { + t.Error("cfg.Batch.Registry is nil, want the assembled tool registry") + } + if cfg.Stream == nil { + t.Error("cfg.Stream is nil, want a stream fn derived from the provider") + } +} + +// TestFreshSessionPersists starts a fresh session, appends a turn to the context, +// persists it, and confirms it round-trips back through the store (FR-16 persist). +func TestFreshSessionPersists(t *testing.T) { + store := newTestStore(t) + s, history, err := newRunSessionWithStore(store, Options{Model: "m", ProviderName: "p"}) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + if history != nil { + t.Fatalf("fresh session history = %v, want nil", history) + } + + s.agentCtx.Messages = append(s.agentCtx.Messages, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("yo")}}, + ) + if err := s.persist(); err != nil { + t.Fatalf("persist: %v", err) + } + if s.persisted != 2 { + t.Errorf("persisted = %d, want 2", s.persisted) + } + + _, msgs, err := store.Load(s.header.ID) + if err != nil { + t.Fatalf("Load: %v", err) + } + if len(msgs) != 2 { + t.Fatalf("persisted messages = %d, want 2", len(msgs)) + } + + // A second persist with no new messages is a no-op. + before := s.curLeaf + if err := s.persist(); err != nil { + t.Fatalf("persist (no-op): %v", err) + } + if s.curLeaf != before { + t.Errorf("curLeaf changed on no-op persist: %q -> %q", before, s.curLeaf) + } +} + +// TestPersistAfterCompaction reproduces the crash where an automatic compaction +// shrinks agentCtx.Messages below the persisted cursor: an incremental +// Messages[persisted:] would panic with a slice-bounds error. persist() must +// instead re-save the flattened context and reset the cursor to the new length. +func TestPersistAfterCompaction(t *testing.T) { + store := newTestStore(t) + s, _, err := newRunSessionWithStore(store, Options{Model: "m", ProviderName: "p"}) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + + // Persist a few turns so the cursor advances past what compaction will keep. + for i := 0; i < 4; i++ { + s.agentCtx.Messages = append(s.agentCtx.Messages, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("q")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("a")}}, + ) + } + if err := s.persist(); err != nil { + t.Fatalf("persist: %v", err) + } + if s.persisted != 8 { + t.Fatalf("persisted = %d, want 8 before compaction", s.persisted) + } + + // Simulate the run loop compacting: Messages is rewritten to a shorter + // summary + tail (here just a 2-message tail), and the loop signalled it via + // compactionMsg (which sets s.compacted). + s.agentCtx.Messages = agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("recent q")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("recent a")}}, + } + s.compacted = true + + if err := s.persist(); err != nil { + t.Fatalf("persist after compaction: %v", err) + } + if s.compacted { + t.Error("compacted flag should be cleared after persist") + } + if s.persisted != 2 { + t.Errorf("persisted = %d, want 2 (the compacted length)", s.persisted) + } + + _, msgs, err := store.Load(s.header.ID) + if err != nil { + t.Fatalf("Load: %v", err) + } + if len(msgs) != 2 { + t.Fatalf("persisted messages = %d, want 2 (flattened compacted context)", len(msgs)) + } +} diff --git a/pigo/internal/cli/tui/slash.go b/pigo/internal/cli/tui/slash.go new file mode 100644 index 0000000..8eadef4 --- /dev/null +++ b/pigo/internal/cli/tui/slash.go @@ -0,0 +1,206 @@ +// This file implements slash-commands and their autocomplete popup for the +// full-screen TUI (US-008, FR-15). It is the TUI counterpart to the REPL's +// slash handling (internal/cli/repl/repl.go): both front-ends consult the SAME +// shared registry assembled by internal/cli/prompts.BuildSlashRegistry (#383), +// so /model, /help, user-declared templates (~/.pigo/{commands,prompts}), +// config/CLI prompt templates, plugin commands and ~/.agents/skills /skill-name +// commands are identical across the two surfaces. +// +// tui deliberately imports prompts/runtime/cli (the shared lower layers), never +// repl: prompts sits below both front-ends, so there is no import cycle. +// +// The autocomplete popup (slashMenu) activates while the input buffer is a +// "/name" being typed (a leading "/" with no whitespace yet). It filters the +// registry by the typed prefix, is navigated with the arrow keys, completed with +// Tab, and run with Enter — the model intercepts those keys before delegating to +// the textarea (see model.handleKey). +package tui + +import ( + "fmt" + "os" + "strings" + + "github.com/smallnest/pigo/internal/cli" + "github.com/smallnest/pigo/internal/cli/prompts" + "github.com/smallnest/pigo/internal/runtime" +) + +// maxMenuRows caps how many candidate rows the popup shows at once; a longer +// filtered list scrolls a window around the selection so the overlay stays a few +// lines tall regardless of how many commands are registered. +const maxMenuRows = 8 + +// newSlashRegistry assembles the shared slash-command registry for the TUI the +// same way the REPL does: built-ins seeded from runtime, the live-state /model +// and /help commands bound to live (so a /model switch mutates the very config +// the run loop reads), user/plugin/skill/template commands from disk. A load +// error is non-fatal — BuildSlashRegistry still returns a registry with the +// built-ins, so the TUI stays usable and the failure is surfaced on stderr. +func newSlashRegistry(opts Options, live *cli.LiveConfig) *runtime.SlashRegistry { + reg, err := prompts.BuildSlashRegistry(live, opts.Skills, opts.Plugins, prompts.PromptTemplateSources{ + Settings: opts.ConfigPrompts, + CLI: opts.CliPrompts, + Disable: opts.NoPromptTemplates, + }) + if err != nil { + fmt.Fprintf(os.Stderr, "pigo: slash-commands: %v\n", err) + } + return reg +} + +// slashMenu is the autocomplete popup state. It holds the candidates matching +// the current "/prefix" and the highlighted row; it is inactive (rendered as +// nothing) whenever the buffer is not a slash-command being typed or no command +// matches the prefix. +type slashMenu struct { + theme Theme + active bool + filtered []runtime.SlashCommand + selected int +} + +// newSlashMenu builds an inactive menu bound to the theme used for its rows. +func newSlashMenu(theme Theme) slashMenu { return slashMenu{theme: theme} } + +// slashToken reports whether buffer is a slash-command name still being typed +// and returns the text after the leading "/". It is true only for a leading "/" +// with no whitespace yet: once the user types a space the name is complete and +// the buffer has moved on to arguments, so name-completion stops. +func slashToken(buffer string) (token string, ok bool) { + trimmed := strings.TrimLeft(buffer, " \t") + if !strings.HasPrefix(trimmed, "/") { + return "", false + } + rest := trimmed[1:] + if strings.ContainsAny(rest, " \t\n") { + return "", false + } + return rest, true +} + +// refresh recomputes the menu from the current buffer and registry. It activates +// only when the buffer is a "/name" prefix that matches at least one command; +// otherwise it deactivates and clears its candidates. The selection is clamped +// so it stays in range as the filtered set shrinks. +func (mn *slashMenu) refresh(buffer string, reg *runtime.SlashRegistry) { + token, ok := slashToken(buffer) + if !ok || reg == nil { + mn.close() + return + } + var out []runtime.SlashCommand + for _, c := range reg.List() { + if strings.HasPrefix(c.Name, token) { + out = append(out, c) + } + } + mn.filtered = out + mn.active = len(out) > 0 + if mn.selected >= len(out) || mn.selected < 0 { + mn.selected = 0 + } +} + +// rows reports how many terminal rows the popup occupies when rendered, so the +// model can reserve that space above the input line during relayout. It is zero +// while inactive and otherwise the visible window height (min of the candidate +// count and maxMenuRows). +func (mn slashMenu) rows() int { + if !mn.active || len(mn.filtered) == 0 { + return 0 + } + if len(mn.filtered) > maxMenuRows { + return maxMenuRows + } + return len(mn.filtered) +} + +// close deactivates the menu and drops its candidates. +func (mn *slashMenu) close() { + mn.active = false + mn.filtered = nil + mn.selected = 0 +} + +// moveUp / moveDown cycle the highlighted candidate, wrapping at the ends so +// arrow navigation is continuous. +func (mn *slashMenu) moveUp() { + if len(mn.filtered) == 0 { + return + } + mn.selected-- + if mn.selected < 0 { + mn.selected = len(mn.filtered) - 1 + } +} + +func (mn *slashMenu) moveDown() { + if len(mn.filtered) == 0 { + return + } + mn.selected++ + if mn.selected >= len(mn.filtered) { + mn.selected = 0 + } +} + +// current returns the highlighted candidate, or ok=false when the menu is +// inactive / empty. +func (mn slashMenu) current() (runtime.SlashCommand, bool) { + if !mn.active || mn.selected < 0 || mn.selected >= len(mn.filtered) { + return runtime.SlashCommand{}, false + } + return mn.filtered[mn.selected], true +} + +// view renders the popup as a block of up to maxMenuRows lines, the highlighted +// row marked with a "›" caret and accented. Each row is "/name description", +// truncated to the width so it never wraps. Returns "" when inactive so the +// model omits the overlay entirely (and its row) while idle. +func (mn slashMenu) view(width int) string { + if !mn.active || len(mn.filtered) == 0 { + return "" + } + start, end := mn.window() + rowWidth := width - 2 // reserve the caret / indent column + if rowWidth < 1 { + rowWidth = width + } + var b strings.Builder + for i := start; i < end; i++ { + c := mn.filtered[i] + line := "/" + c.Name + if c.Description != "" { + line += " " + c.Description + } + line = TruncateToWidth(line, rowWidth) + if i == mn.selected { + b.WriteString(mn.theme.Accent.Render("› " + line)) + } else { + b.WriteString(mn.theme.System.Render(" " + line)) + } + if i < end-1 { + b.WriteByte('\n') + } + } + return b.String() +} + +// window returns the [start,end) slice of filtered candidates to display, +// scrolled to keep the selection visible when the list is taller than +// maxMenuRows. +func (mn slashMenu) window() (int, int) { + n := len(mn.filtered) + if n <= maxMenuRows { + return 0, n + } + start := mn.selected - maxMenuRows + 1 + if start < 0 { + start = 0 + } + if start > n-maxMenuRows { + start = n - maxMenuRows + } + return start, start + maxMenuRows +} diff --git a/pigo/internal/cli/tui/slash_test.go b/pigo/internal/cli/tui/slash_test.go new file mode 100644 index 0000000..6b32369 --- /dev/null +++ b/pigo/internal/cli/tui/slash_test.go @@ -0,0 +1,216 @@ +package tui + +import ( + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/runtime" +) + +// typeInto feeds each rune of s to the model as a key press, returning the +// evolved model. It mirrors how a terminal delivers typed characters (Code + +// Text), including the leading "/" of a slash-command. +func typeInto(t *testing.T, m tea.Model, s string) tea.Model { + t.Helper() + for _, r := range s { + m, _ = m.Update(tea.KeyPressMsg{Code: r, Text: string(r)}) + } + return m +} + +// menuNames returns the "/name" of every candidate currently in the popup. +func menuNames(m Model) []string { + out := make([]string, len(m.menu.filtered)) + for i, c := range m.menu.filtered { + out[i] = "/" + c.Name + } + return out +} + +func containsAll(hay []string, needles ...string) bool { + set := make(map[string]bool, len(hay)) + for _, h := range hay { + set[h] = true + } + for _, n := range needles { + if !set[n] { + return false + } + } + return true +} + +// TestSlashMenuOpensOnSlash verifies that typing a bare "/" opens the popup with +// the built-in commands present (/model, /help among them). +func TestSlashMenuOpensOnSlash(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/").(Model) + if !m.menu.active { + t.Fatalf("menu should be active after typing '/'") + } + names := menuNames(m) + if !containsAll(names, "/model", "/help") { + t.Errorf("candidate set %v missing /model or /help", names) + } +} + +// TestSlashMenuFiltersByPrefix verifies the popup narrows to the typed prefix: +// "/mo" keeps /model and /models but drops /help. +func TestSlashMenuFiltersByPrefix(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/mo").(Model) + if !m.menu.active { + t.Fatalf("menu should be active for '/mo'") + } + names := menuNames(m) + if !containsAll(names, "/model", "/models") { + t.Errorf("candidate set %v missing /model or /models", names) + } + for _, n := range names { + if !strings.HasPrefix(n, "/mo") { + t.Errorf("candidate %q does not match prefix /mo (set %v)", n, names) + } + } +} + +// TestSlashMenuClosesOnSpace verifies name-completion stops once the buffer moves +// on to arguments (a space after the command name closes the popup). +func TestSlashMenuClosesOnSpace(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/model ").(Model) + if m.menu.active { + t.Errorf("menu should close once the command name is complete (buffer %q)", m.input.Value()) + } +} + +// TestSlashMenuNavigation verifies arrow keys move the highlighted candidate and +// wrap at the ends. +func TestSlashMenuNavigation(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/").(Model) + n := len(m.menu.filtered) + if n < 2 { + t.Fatalf("need at least two candidates to test navigation, got %d", n) + } + next, _ := m.Update(tea.KeyPressMsg{Code: tea.KeyDown}) + if got := next.(Model).menu.selected; got != 1 { + t.Errorf("after Down, selected = %d, want 1", got) + } + // Up from index 0 wraps to the last candidate. + back, _ := m.Update(tea.KeyPressMsg{Code: tea.KeyUp}) + if got := back.(Model).menu.selected; got != n-1 { + t.Errorf("after Up from 0, selected = %d, want %d (wrap)", got, n-1) + } +} + +// TestSlashTabCompletes verifies Tab fills the buffer with the highlighted +// command and closes the popup (ready for arguments). +func TestSlashTabCompletes(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/hel").(Model) + got, _ := m.Update(tea.KeyPressMsg{Code: tea.KeyTab}) + gm := got.(Model) + if gm.input.Value() != "/help " { + t.Errorf("after Tab, buffer = %q, want %q", gm.input.Value(), "/help ") + } + if gm.menu.active { + t.Errorf("menu should close after Tab completion") + } +} + +// TestSlashHelpExecutesIntoTranscript verifies executing a built-in action +// command (/help) renders its output into the transcript as a system block — +// listing the available commands — without requiring a live provider. +func TestSlashHelpExecutesIntoTranscript(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/help").(Model) + got, cmd := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + gm := got.(Model) + // No run is started for an action command. + if cmd != nil { + if msg := cmd(); msg != nil { + if _, isQuit := msg.(tea.QuitMsg); isQuit { + t.Fatalf("/help should not quit") + } + } + } + if gm.running { + t.Errorf("/help is an action command; model should stay idle") + } + joined := strings.Join(blockTexts(gm.transcript), "\n") + if !strings.Contains(joined, "/help") || !strings.Contains(joined, "/model") { + t.Errorf("/help output should list commands (/help, /model); transcript:\n%s", joined) + } + if gm.input.Value() != "" { + t.Errorf("after executing /help, input = %q, want cleared", gm.input.Value()) + } +} + +// TestSlashUnknownCommandReported verifies an unknown "/name" surfaces the +// resolver error into the transcript rather than being sent to the agent. +func TestSlashUnknownCommandReported(t *testing.T) { + m := typeInto(t, NewModel(Options{}), "/definitelynotacommand").(Model) + // The popup filters to nothing, so it is inactive; Enter routes through submit. + got, _ := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + gm := got.(Model) + if gm.running { + t.Errorf("an unknown command must not start a run") + } + joined := strings.Join(blockTexts(gm.transcript), "\n") + if !strings.Contains(joined, "unknown command") { + t.Errorf("expected an unknown-command notice in transcript, got:\n%s", joined) + } +} + +// TestSlashMenuRendersAboveInput verifies the popup appears in the View while +// active, so the candidate list is visible above the input line. +func TestSlashMenuRendersAboveInput(t *testing.T) { + m := NewModel(Options{}) + next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + m = typeInto(t, next, "/mo").(Model) + view := m.View() + if !strings.Contains(view.Content, "/model") { + t.Errorf("active popup should render /model in the view content") + } +} + +// TestSlashPromptCommandStartsRun verifies a prompt (Expand) command's expanded +// text is fed to the run seam, not shown as a bare status. A stub startRunFn +// stands in for the live provider. +func TestSlashPromptCommandStartsRun(t *testing.T) { + m := NewModel(Options{}) + // Inject a user prompt command directly into the registry. + m.slash.AddUser(runtime.SlashCommand{ + Name: "greet", + Expand: func(args string) string { return "hello " + args }, + }) + var ran string + m.startRunFn = func(prompt string) (chan tea.Msg, tea.Cmd) { + ran = prompt + ch := make(chan tea.Msg, 1) + return ch, func() tea.Msg { return nil } + } + got, _ := m.runSlash("/greet world") + gm := got.(Model) + if ran != "hello world" { + t.Errorf("prompt command should start a run with expanded text; got %q", ran) + } + if !gm.running { + t.Errorf("model should be running after a prompt command") + } +} + +// TestSlashExitQuits verifies that /exit and /quit typed in the TUI input box +// terminate the program (tea.Quit + quitting flag), mirroring the REPL loop +// which intercepts them before slash resolution. +func TestSlashExitQuits(t *testing.T) { + for _, cmd := range []string{"/exit", "/quit"} { + got, teaCmd := NewModel(Options{}).runSlash(cmd) + gm := got.(Model) + if !gm.quitting { + t.Errorf("%s: model should be marked quitting", cmd) + } + if teaCmd == nil { + t.Fatalf("%s: expected a tea.Quit command, got nil", cmd) + } + if _, isQuit := teaCmd().(tea.QuitMsg); !isQuit { + t.Errorf("%s: cmd should be tea.Quit", cmd) + } + } +} diff --git a/pigo/internal/cli/tui/spinner.go b/pigo/internal/cli/tui/spinner.go new file mode 100644 index 0000000..8d58846 --- /dev/null +++ b/pigo/internal/cli/tui/spinner.go @@ -0,0 +1,185 @@ +package tui + +import ( + "fmt" + "math/rand" + "strings" + "time" +) + +// This file implements the "working" spinner shown while an agent run is in +// flight, mirroring Claude Code's animated status line: a cycling asterisk +// glyph, a whimsical present-progressive verb ("Whirring…"), and a live stats +// readout — elapsed wall-clock time, an estimate of streamed output tokens, and +// the configured thinking effort. It renders on the row just above the input +// while running and disappears when the run ends. + +// spinnerTickMsg advances the spinner animation. The model re-issues a tick +// after each frame while a run is in flight and lets the tick lapse once the run +// ends, so the animation stops without a running goroutine. +type spinnerTickMsg time.Time + +// spinnerInterval is the frame cadence. ~120ms is brisk enough to read as motion +// without churning the render loop. +const spinnerInterval = 120 * time.Millisecond + +// verbRerollFrames re-picks the verb roughly every this many frames (~5s) so a +// long run cycles through several verbs the way Claude Code does. +const verbRerollFrames = 40 + +// spinnerFrames is the asterisk animation cycled one glyph per tick. The glyphs +// grow from a dim dot to a full star and back, reading as a pulsing sparkle. +var spinnerFrames = []string{"·", "✢", "✳", "∗", "✺", "✻", "✽", "✻", "✺", "∗", "✳", "✢"} + +// spinner is the animated working indicator. It is a plain value held by the +// Model: begin() arms it at run start, advance() steps the frame on each tick, +// addTokens() grows the streamed-token estimate, and view() renders the line. +type spinner struct { + theme Theme + running bool + frame int + verb string + start time.Time + chars int // runes streamed this run (the token estimate divides this) + thinking string // thinking-effort label, e.g. "medium"; "" hides that stat + pinned string // when set, overrides the random verb and stops re-rolling +} + +// newSpinner builds an idle spinner bound to the theme. +func newSpinner(theme Theme) spinner { + return spinner{theme: theme} +} + +// begin arms the spinner for a fresh run: it records the start time, picks the +// first verb, resets the frame and token estimate, and stores the thinking-effort +// label to show in the stats. +func (s *spinner) begin(now time.Time, thinking string) { + s.running = true + s.frame = 0 + s.start = now + s.chars = 0 + s.thinking = thinking + s.verb = randomVerb() + s.pinned = "" +} + +// pin fixes the spinner label to a specific phrase (e.g. "Compacting +// conversation") and stops verb re-rolling until unpin, so a long-running phase +// reads as one steady message rather than cycling words. +func (s *spinner) pin(label string) { s.pinned = label } + +// unpin restores the normal cycling verb after a pinned phase ends. +func (s *spinner) unpin() { s.pinned = "" } + +// stop parks the spinner when a run ends so view() renders nothing. +func (s *spinner) stop() { s.running = false } + +// advance steps the animation one frame and periodically re-rolls the verb so a +// long run does not sit on one word. +func (s *spinner) advance() { + s.frame++ + if s.pinned == "" && s.frame%verbRerollFrames == 0 { + s.verb = randomVerb() + } +} + +// addTokens folds a streamed text delta into the running output-token estimate. +// The count is approximate (≈4 chars per token) — enough for a live spinner +// readout, not billing. +func (s *spinner) addTokens(delta string) { + s.chars += len([]rune(delta)) +} + +// view renders the spinner line, e.g. "✻ Whirring… (1m 54s · ↓ 242 tokens · +// medium effort)". It returns "" when not running or before a width is known. +// The glyph and verb take the accent color; the parenthetical stats are dim. +func (s spinner) view(width int) string { + if !s.running || width <= 0 { + return "" + } + glyph := spinnerFrames[s.frame%len(spinnerFrames)] + verb := s.verb + if s.pinned != "" { + verb = s.pinned + } + head := s.theme.Spinner.Render(glyph + " " + verb + "…") + + var stats strings.Builder + fmt.Fprintf(&stats, "%s", formatElapsed(time.Since(s.start))) + if tokens := s.chars / 4; tokens > 0 { + fmt.Fprintf(&stats, " · ↓ %s tokens", humanizeInt(tokens)) + } + if s.thinking != "" { + fmt.Fprintf(&stats, " · %s effort", s.thinking) + } + line := head + " " + s.theme.System.Render("("+stats.String()+")") + return TruncateToWidth(line, width) +} + +// randomVerb picks one of the built-in present-progressive verbs. +func randomVerb() string { + return spinnerVerbs[rand.Intn(len(spinnerVerbs))] +} + +// formatElapsed renders a duration compactly: "42s", "1m 54s", or "1h 2m". +func formatElapsed(d time.Duration) string { + if d < 0 { + d = 0 + } + secs := int(d.Seconds()) + if secs < 60 { + return fmt.Sprintf("%ds", secs) + } + mins := secs / 60 + secs %= 60 + if mins < 60 { + return fmt.Sprintf("%dm %ds", mins, secs) + } + hours := mins / 60 + mins %= 60 + return fmt.Sprintf("%dh %dm", hours, mins) +} + +// spinnerVerbs is Claude Code's 185 built-in spinner verbs (present-progressive +// flavor words shown while working). Sourced from the community catalog at +// github.com/wynandw87/claude-code-spinner-verbs. +var spinnerVerbs = []string{ + "Accomplishing", "Actioning", "Actualizing", "Architecting", "Baking", + "Beaming", "Beboppin'", "Befuddling", "Billowing", "Blanching", + "Bloviating", "Boogieing", "Boondoggling", "Booping", "Bootstrapping", + "Brewing", "Burrowing", "Calculating", "Canoodling", "Caramelizing", + "Cascading", "Catapulting", "Cerebrating", "Channeling", "Channelling", + "Choreographing", "Churning", "Clauding", "Coalescing", "Cogitating", + "Combobulating", "Composing", "Computing", "Concocting", "Considering", + "Contemplating", "Cooking", "Crafting", "Creating", "Crunching", + "Crystallizing", "Cultivating", "Deciphering", "Deliberating", "Determining", + "Dilly-dallying", "Discombobulating", "Doing", "Doodling", "Drizzling", + "Ebbing", "Effecting", "Elucidating", "Embellishing", "Enchanting", + "Envisioning", "Evaporating", "Fermenting", "Fiddle-faddling", "Finagling", + "Flambeing", "Flibbertigibbeting", "Flowing", "Flummoxing", "Fluttering", + "Forging", "Forming", "Frolicking", "Frosting", "Gallivanting", + "Galloping", "Garnishing", "Generating", "Germinating", "Gitifying", + "Grooving", "Gusting", "Harmonizing", "Hashing", "Hatching", + "Herding", "Honking", "Hullaballooing", "Hyperspacing", "Ideating", + "Imagining", "Improvising", "Incubating", "Inferring", "Infusing", + "Ionizing", "Jitterbugging", "Julienning", "Kneading", "Leavening", + "Levitating", "Lollygagging", "Manifesting", "Marinating", "Meandering", + "Metamorphosing", "Misting", "Moonwalking", "Moseying", "Mulling", + "Mustering", "Musing", "Nebulizing", "Nesting", "Newspapering", + "Noodling", "Nucleating", "Orbiting", "Orchestrating", "Osmosing", + "Perambulating", "Percolating", "Perusing", "Philosophising", "Photosynthesizing", + "Pollinating", "Pondering", "Pontificating", "Pouncing", "Precipitating", + "Prestidigitating", "Processing", "Proofing", "Propagating", "Puttering", + "Puzzling", "Quantumizing", "Razzle-dazzling", "Razzmatazzing", "Recombobulating", + "Reticulating", "Roosting", "Ruminating", "Sauteing", "Scampering", + "Schlepping", "Scurrying", "Seasoning", "Shenaniganing", "Shimmying", + "Simmering", "Skedaddling", "Sketching", "Slithering", "Smooshing", + "Sock-hopping", "Spelunking", "Spinning", "Sprouting", "Stewing", + "Sublimating", "Swirling", "Swooping", "Symbioting", "Synthesizing", + "Tempering", "Thinking", "Thundering", "Tinkering", "Tomfoolering", + "Topsy-turvying", "Transfiguring", "Transmuting", "Twisting", "Undulating", + "Unfurling", "Unravelling", "Vibing", "Waddling", "Wandering", + "Warping", "Whatchamacalliting", "Whirlpooling", "Whirring", "Whisking", + "Wibbling", "Working", "Wrangling", "Zesting", "Zigzagging", +} + diff --git a/pigo/internal/cli/tui/spinner_test.go b/pigo/internal/cli/tui/spinner_test.go new file mode 100644 index 0000000..d5810cf --- /dev/null +++ b/pigo/internal/cli/tui/spinner_test.go @@ -0,0 +1,123 @@ +package tui + +import ( + "strings" + "testing" + "time" + + tea "charm.land/bubbletea/v2" +) + +// TestFormatElapsed checks the compact duration formatting across the second, +// minute, and hour ranges. +func TestFormatElapsed(t *testing.T) { + cases := []struct { + d time.Duration + want string + }{ + {5 * time.Second, "5s"}, + {59 * time.Second, "59s"}, + {114 * time.Second, "1m 54s"}, + {60 * time.Minute, "1h 0m"}, + {62 * time.Minute, "1h 2m"}, + {-3 * time.Second, "0s"}, + } + for _, c := range cases { + if got := formatElapsed(c.d); got != c.want { + t.Errorf("formatElapsed(%s) = %q, want %q", c.d, got, c.want) + } + } +} + +// TestSpinnerViewStats verifies a running spinner renders its verb with an +// ellipsis and the elapsed/token/effort stats, and that a stopped spinner +// renders nothing. +func TestSpinnerViewStats(t *testing.T) { + s := newSpinner(DefaultTheme()) + s.begin(time.Now().Add(-114*time.Second), "medium") + s.chars = 968 // 968/4 = 242 estimated tokens + + view := stripANSI(s.view(120)) + if !strings.Contains(view, s.verb+"…") { + t.Errorf("view %q should contain the verb with an ellipsis", view) + } + for _, want := range []string{"1m 54s", "↓ 242 tokens", "medium effort"} { + if !strings.Contains(view, want) { + t.Errorf("view %q missing stat %q", view, want) + } + } + + s.stop() + if got := s.view(120); got != "" { + t.Errorf("stopped spinner should render nothing, got %q", got) + } +} + +// TestSpinnerPinOverridesVerb verifies a pinned label replaces the random verb +// and survives verb re-rolls, and that unpin restores the cycling verb. +func TestSpinnerPinOverridesVerb(t *testing.T) { + s := newSpinner(DefaultTheme()) + s.begin(time.Now(), "") + s.pin("Compacting conversation") + + // Advance well past the re-roll interval: a pinned label must not change. + for i := 0; i < verbRerollFrames*2; i++ { + s.advance() + } + view := stripANSI(s.view(120)) + if !strings.Contains(view, "Compacting conversation…") { + t.Errorf("pinned spinner view %q should show the pinned label", view) + } + + s.unpin() + if got := stripANSI(s.view(120)); strings.Contains(got, "Compacting conversation") { + t.Errorf("after unpin, view %q should not show the pinned label", got) + } +} +// tokens stream and the effort stat is hidden with no thinking level. +func TestSpinnerViewOmitsEmptyStats(t *testing.T) { + s := newSpinner(DefaultTheme()) + s.begin(time.Now(), "") + + view := stripANSI(s.view(120)) + if strings.Contains(view, "tokens") { + t.Errorf("view %q should not show a token stat before any deltas", view) + } + if strings.Contains(view, "effort") { + t.Errorf("view %q should not show an effort stat with no thinking level", view) + } +} + +// TestSpinnerAdvanceRerollsVerb verifies the animation frame advances and the +// verb is re-picked on the reroll cadence. +func TestSpinnerAdvanceRerollsVerb(t *testing.T) { + s := newSpinner(DefaultTheme()) + s.begin(time.Now(), "") + if s.frame != 0 { + t.Fatalf("fresh spinner frame = %d, want 0", s.frame) + } + for i := 0; i < verbRerollFrames; i++ { + s.advance() + } + if s.frame != verbRerollFrames { + t.Errorf("frame after %d advances = %d", verbRerollFrames, s.frame) + } +} + +// TestModelRunningShowsSpinnerRow verifies that while a run is in flight the +// spinner occupies its own row above the input and the shell still fills exactly +// the terminal height (relayout shrinks the transcript by the spinner row). +func TestModelRunningShowsSpinnerRow(t *testing.T) { + m := apply(t, NewModel(Options{Model: "test-model"}), tea.WindowSizeMsg{Width: 60, Height: 10}) + m.running = true + m.spinner.begin(time.Now(), "medium") + m.relayout() + + view := m.renderContent() + if got := strings.Count(view, "\n"); got != 9 { + t.Errorf("running newline count = %d, want 9 (10 rows)", got) + } + if !strings.Contains(stripANSI(view), m.spinner.verb+"…") { + t.Errorf("running view should contain the spinner verb, got:\n%s", stripANSI(view)) + } +} diff --git a/pigo/internal/cli/tui/status_test.go b/pigo/internal/cli/tui/status_test.go new file mode 100644 index 0000000..89a836d --- /dev/null +++ b/pigo/internal/cli/tui/status_test.go @@ -0,0 +1,219 @@ +package tui + +import ( + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// typeCommand feeds "/name" into the model and presses Enter, mirroring how a +// user runs a slash command from the composer (the popup is open at Enter, so it +// routes through submitSlashSelected, exactly like the REPL path). +func typeCommand(t *testing.T, m Model, cmd string) Model { + t.Helper() + m = typeInto(t, m, cmd).(Model) + got, c := m.Update(tea.KeyPressMsg{Code: tea.KeyEnter}) + if c != nil { + if msg := c(); msg != nil { + if _, isQuit := msg.(tea.QuitMsg); isQuit { + t.Fatalf("%s should not quit", cmd) + } + } + } + return got.(Model) +} + +// TestStatusWithSessionRendersSections drives /status on a session-bound model +// and asserts every report section appears in the transcript, with the model +// staying idle (no run is started). +func TestStatusWithSessionRendersSections(t *testing.T) { + store := newTestStore(t) + s, _, err := newRunSessionWithStore(store, Options{ + Model: "status-model", + ProviderName: "status-provider", + ThinkingLevel: agentcore.ThinkingLevel("low"), + }) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + // The /status environment section keys off the launch directory; keep it + // deterministic rather than depending on the test's cwd. + s.cwd = "/tmp/tui-status" + + m := NewModel(Options{}).withSession(s, nil) + m = typeCommand(t, m, "/status") + + if m.running { + t.Error("/status is an action command; model should stay idle") + } + if m.input.Value() != "" { + t.Errorf("after executing /status, input = %q, want cleared", m.input.Value()) + } + joined := strings.Join(blockTexts(m.transcript), "\n") + for _, want := range []string{ + "runtime config:", + "model: status-model", + "provider: status-provider", + "context:", + "project & environment:", + "cwd: /tmp/tui-status", + "credentials & connectivity:", + "telemetry:", + "no telemetry yet", + } { + if !strings.Contains(joined, want) { + t.Errorf("/status output missing %q; transcript:\n%s", want, joined) + } + } +} + +// TestStatusWithoutSessionNotice verifies a session-less model reports the +// unavailable notice rather than panicking on nil collaborators, and still +// clears the input. +func TestStatusWithoutSessionNotice(t *testing.T) { + m := NewModel(Options{}) + m = typeCommand(t, m, "/status") + + if m.running { + t.Error("/status must not start a run on a session-less model") + } + if m.input.Value() != "" { + t.Errorf("after executing /status, input = %q, want cleared", m.input.Value()) + } + joined := strings.Join(blockTexts(m.transcript), "\n") + if !strings.Contains(joined, "status unavailable: no active session") { + t.Errorf("expected an unavailable notice in transcript, got:\n%s", joined) + } +} + +// TestSessionCommandRendersSummary drives /session on a session-bound model and +// asserts the summary lines (session id, message count, tokens, model/provider, +// compactions) match the REPL's /session format. +func TestSessionCommandRendersSummary(t *testing.T) { + store := newTestStore(t) + s, _, err := newRunSessionWithStore(store, Options{ + Model: "session-model", + ProviderName: "session-provider", + }) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + // Simulate two user/assistant turns in the live context (unsaved messages are + // counted too, mirroring the REPL's in-memory source of truth). + s.agentCtx.Messages = agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("q1")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("a1")}}, + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("q2")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("a2")}}, + } + + m := NewModel(Options{}).withSession(s, nil) + m = typeCommand(t, m, "/session") + + if m.running { + t.Error("/session is an action command; model should stay idle") + } + joined := strings.Join(blockTexts(m.transcript), "\n") + for _, want := range []string{ + "session: " + s.header.ID, + "messages: 4", + "tokens (est):", + "model: session-model (provider: session-provider)", + "compactions: 0", + } { + if !strings.Contains(joined, want) { + t.Errorf("/session output missing %q; transcript:\n%s", want, joined) + } + } +} + +// TestSessionWithoutSessionNotice verifies a session-less model reports the +// unavailable notice for /session. +func TestSessionWithoutSessionNotice(t *testing.T) { + m := NewModel(Options{}) + m = typeCommand(t, m, "/session") + + joined := strings.Join(blockTexts(m.transcript), "\n") + if !strings.Contains(joined, "session unavailable: no active session") { + t.Errorf("expected an unavailable notice in transcript, got:\n%s", joined) + } +} + +// TestTelemetryFoldingFeedsStatus verifies a telemetryMsg from the bridge is +// folded into the session's telemetry holder, so /status renders the cumulative +// and last-run telemetry blocks instead of "no telemetry yet". +func TestTelemetryFoldingFeedsStatus(t *testing.T) { + store := newTestStore(t) + s, _, err := newRunSessionWithStore(store, Options{Model: "m", ProviderName: "p"}) + if err != nil { + t.Fatalf("newRunSessionWithStore: %v", err) + } + if s.telemetry == nil { + t.Fatal("session telemetry holder should be initialized") + } + + m := NewModel(Options{}).withSession(s, nil) + // A telemetryMsg is bridged while a run is pumping; set up that state so + // Update keeps the pump running (pumpNext) exactly as during a real run. + m.running = true + m.runCh = make(chan tea.Msg) + next, cmd := m.Update(telemetryMsg{ev: agentcore.TelemetryEvent{ + Turns: 3, + TruncationCount: 1, + CompactionCount: 1, + ContextTokens: 1000, + ContextWindow: 200000, + ContextUtilization: 0.005, + }}) + if cmd == nil { + t.Fatal("telemetryMsg should return a pump cmd while a run is in flight") + } + m = next.(Model) + + // The event must be retained on the session's holder (not just the status + // bar), otherwise /status could not render the telemetry section. + if !s.telemetry.HasTelemetry() { + t.Fatal("telemetry event should be folded into the session holder") + } + if s.telemetry.CumulativeTurns() != 3 { + t.Errorf("CumulativeTurns = %d, want 3", s.telemetry.CumulativeTurns()) + } + + // The run has since ended, returning the model to idle so the user can + // issue /status (key presses are dropped while a run is in flight). + m.running = false + m.runCh = nil + + m = typeCommand(t, m, "/status") + joined := strings.Join(blockTexts(m.transcript), "\n") + for _, want := range []string{ + "telemetry:", + "since session start:", + "turns: 3", + "truncations: 1", + "last run:", + } { + if !strings.Contains(joined, want) { + t.Errorf("/status output missing %q after telemetry fold; transcript:\n%s", want, joined) + } + } + if strings.Contains(joined, "no telemetry yet") { + t.Error("/status should render real telemetry after a fold, not 'no telemetry yet'") + } +} + +// TestStatusNotInterceptedForStatusFoo verifies "/statusfoo" is NOT intercepted +// as /status (mirroring the REPL's guard), so it resolves as an unknown command +// rather than rendering the status report. +func TestStatusNotInterceptedForStatusFoo(t *testing.T) { + m := NewModel(Options{}) + m = typeCommand(t, m, "/statusfoo") + + joined := strings.Join(blockTexts(m.transcript), "\n") + if strings.Contains(joined, "runtime config:") { + t.Errorf("/statusfoo must not run the status command; transcript:\n%s", joined) + } +} diff --git a/pigo/internal/cli/tui/statusbar.go b/pigo/internal/cli/tui/statusbar.go new file mode 100644 index 0000000..533601c --- /dev/null +++ b/pigo/internal/cli/tui/statusbar.go @@ -0,0 +1,397 @@ +package tui + +import ( + "fmt" + "os" + "strings" + + "charm.land/lipgloss/v2" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// statusBar renders the persistent bottom line described in the SPEC (US-003, +// Section 5.1): model name, thinking level, cwd (with $HOME abbreviated to ~), +// git branch + dirty/ahead markers, context-usage %, and the current task text. +// It holds no styling of its own beyond the Theme's StatusBar style; all width +// fitting is done against ui.Width so CJK/emoji count as two columns. +// +// The component is a plain value: Update-side code copies it into the Model, +// mutates the exported-to-package snapshot fields via the setters, and calls +// Render(width) from View. It never performs I/O — the git probe lives in +// gitinfo.go and feeds it via SetGit. +type statusBar struct { + theme Theme + + // Static-ish config sourced from Options. + model string + thinking string + + // cwd is the launch directory with $HOME already abbreviated to "~". + cwd string + + // git is the latest probe result; rendered only when git.ok is true. + git gitInfoMsg + + // contextPct is the latest context-window utilization in percent [0,100], + // derived from telemetryMsg (ContextUtilization * 100). -1 means unknown, so + // the segment is hidden until the first telemetry arrives. + contextPct int + + // tokens is the most recently observed context-token count (ContextTokens), + // shown alongside the percentage. 0 means unknown/not yet reported. + tokens int + + // task is the current activity text (e.g. the running tool or turn state). + task string +} + +// newStatusBar builds a status bar from the theme, resolved Options, and the +// launch directory. contextPct starts at -1 (unknown) so the token segment stays +// hidden until telemetry arrives. +func newStatusBar(theme Theme, opts Options, cwd string) statusBar { + return statusBar{ + theme: theme, + model: opts.Model, + thinking: string(opts.ThinkingLevel), + cwd: abbreviateHome(cwd), + contextPct: -1, + } +} + +// SetGit stores the latest git probe result. +func (s *statusBar) SetGit(g gitInfoMsg) { s.git = g } + +// SetModel updates the displayed model name after a /model switch. +func (s *statusBar) SetModel(model string) { s.model = model } + +// SetThinking updates the displayed reasoning-effort level after a /think switch. +func (s *statusBar) SetThinking(level string) { s.thinking = level } + +// SetTelemetry updates the context-usage percentage from a telemetry event. +// A zero/unknown window (ContextWindow == 0) leaves the segment hidden. +func (s *statusBar) SetTelemetry(ev telemetryEventView) { + if ev.window <= 0 { + s.contextPct = -1 + s.tokens = 0 + return + } + pct := int(ev.util*100 + 0.5) + if pct < 0 { + pct = 0 + } + if pct > 100 { + pct = 100 + } + s.contextPct = pct + if ev.tokens > 0 { + s.tokens = ev.tokens + } +} + +// SetTask records the current activity text shown at the far right / high +// priority slot of the bar. +func (s *statusBar) SetTask(task string) { s.task = task } + +// telemetryEventView is the minimal projection of agentcore.TelemetryEvent the +// status bar needs, so the caller (model.go) adapts the event rather than this +// file depending on agentcore directly for a two-field read. +type telemetryEventView struct { + util float64 + window int + tokens int +} + +// appName is the badge shown at the far left of the bar. +const appName = "pigo" + +// Glyphs prefixing each segment plus the powerline separator, matching the +// decorated Claude-Code-plugin look. The segment icons are common Unicode; the +// separator (sepArrow) is a powerline glyph in the private-use area that Nerd +// Fonts and most modern terminal fonts render. Each measures one display column +// and ui.Width accounts for it during truncation. +const ( + glyphGit = "⎇" // git branch + glyphDirty = "●" // uncommitted changes + glyphAhead = "⇡" // commits ahead of upstream + glyphModel = "✱" // model name + glyphThink = "✽" // thinking level + glyphCwd = "▸" // working directory + glyphCtx = "◔" // context-window usage + glyphTask = "⏵" // current activity + + sepArrow = "" // filled right arrow — used at a background transition +) + +// Powerline palette (ANSI 256-color cube, so it renders without true-color). +// Every segment is its own colored block. The arrow between two segments is +// drawn in the LEFT block's background color so it reads as that item's color +// spilling into the next; a closing arrow caps the final block back to the bar. +const ( + sbBarBg = "236" // bar background behind the trailing pad + + sbAppFg = "233" // app badge text (dark, on light gray) + sbAppBg = "252" // app badge block (light gray) + + sbGitFg = "231" // git text + sbGitBg = "65" // git block (muted green) + + sbModelFg = "231" // model text + sbModelBg = "97" // model block (muted purple) + + sbThinkFg = "231" // thinking text + sbThinkBg = "60" // thinking block (slate) + + sbCwdFg = "231" // cwd text + sbCwdBg = "67" // cwd block (steel blue) + + sbCtxFg = "236" // context text (dark, on amber) + sbCtxBg = "179" // context block (amber/gold) + + sbTaskFg = "231" // task text + sbTaskBg = "131" // task block (muted terracotta) +) + +// segment is one labelled field of the bar together with its colors and +// truncation priority. bg == "" means the segment sits on the bar background; +// a non-empty bg gives it a filled powerline block. Higher priority survives +// longer when the terminal is too narrow. +type segment struct { + text string + fg string + bg string // "" => bar background + priority int // larger = kept longer under truncation +} + +// Priority order (SPEC: task > model/app > token > git > cwd). Higher is more +// important and dropped/truncated last. +const ( + prioCwd = 0 + prioGit = 1 + prioToken = 2 + prioModel = 3 + prioApp = 3 // the app badge rides at the model tier + prioTask = 4 +) + +// Render lays the bar out to exactly the configured width as a colored powerline +// ribbon. Each segment is a filled block joined to the next by an arrow drawn in +// the left block's background color, and the tail is padded with the bar +// background so the whole row is filled. When the ribbon would exceed the width +// it drops whole segments from lowest to highest priority; if even the single +// highest-priority segment still overflows it hard-truncates that segment's +// text. The rendered row's display width (ui.Width, which ignores ANSI) is +// always exactly width for width > 0; a non-positive width yields the empty +// string. +func (s statusBar) Render(width int) string { + if width <= 0 { + return "" + } + + segs := s.segments() + for len(segs) > 0 { + ribbon, w := renderRibbon(segs) + if w <= width { + return ribbon + barPad(width-w) + } + if len(segs) == 1 { + break + } + segs = dropLowest(segs) + } + + // Even a single highest-priority segment overflows: hard-truncate its text + // onto the bar background (no separators, so width stays bounded). + txt := TruncateToWidth(s.highestText(), width) + base := lipgloss.NewStyle().Foreground(lipgloss.Color(sbAppFg)).Background(lipgloss.Color(sbBarBg)) + return base.Render(txt) + barPad(width-ui.Width(txt)) +} + +// renderRibbon builds the styled powerline string for segs and returns it with +// its visible width (excluding ANSI). The left edge starts at the first +// segment's background; a closing arrow caps any trailing block back to the bar. +func renderRibbon(segs []segment) (string, int) { + resolve := func(bg string) string { + if bg == "" { + return sbBarBg + } + return bg + } + + var b strings.Builder + vis := 0 + for i, seg := range segs { + curBg := resolve(seg.bg) + if i > 0 { + // The separator arrow is filled with the LEFT block's background, so + // it matches the item it flows out of, sitting on the next block's bg. + leftBg := resolve(segs[i-1].bg) + b.WriteString(lipgloss.NewStyle(). + Foreground(lipgloss.Color(leftBg)). + Background(lipgloss.Color(curBg)). + Render(sepArrow)) + vis += ui.Width(sepArrow) + } + content := " " + seg.text + " " + b.WriteString(lipgloss.NewStyle(). + Foreground(lipgloss.Color(seg.fg)). + Background(lipgloss.Color(curBg)). + Render(content)) + vis += ui.Width(content) + } + + // Cap a trailing colored block with an arrow back to the bar background. + if lastBg := resolve(segs[len(segs)-1].bg); lastBg != sbBarBg { + b.WriteString(lipgloss.NewStyle(). + Foreground(lipgloss.Color(lastBg)). + Background(lipgloss.Color(sbBarBg)). + Render(sepArrow)) + vis += ui.Width(sepArrow) + } + return b.String(), vis +} + +// barPad returns n spaces painted with the bar background so the row fills the +// full terminal width. n <= 0 yields the empty string. +func barPad(n int) string { + if n <= 0 { + return "" + } + return lipgloss.NewStyle(). + Background(lipgloss.Color(sbBarBg)). + Render(strings.Repeat(" ", n)) +} + +// segments builds the ordered list of visible segments. Order in the slice is +// the left-to-right display order; priority governs truncation, not position. +func (s statusBar) segments() []segment { + var segs []segment + + // App badge leads the bar as a filled block. + segs = append(segs, segment{text: appName, fg: sbAppFg, bg: sbAppBg, priority: prioApp}) + + if s.git.ok { + segs = append(segs, segment{text: s.gitText(), fg: sbGitFg, bg: sbGitBg, priority: prioGit}) + } + if s.model != "" { + segs = append(segs, segment{text: glyphModel + " " + s.model, fg: sbModelFg, bg: sbModelBg, priority: prioModel}) + } + if s.thinking != "" { + // Thinking rides with the model priority — it is cheap and contextual. + segs = append(segs, segment{text: glyphThink + " " + s.thinking, fg: sbThinkFg, bg: sbThinkBg, priority: prioModel}) + } + if s.cwd != "" { + segs = append(segs, segment{text: glyphCwd + " " + s.cwd, fg: sbCwdFg, bg: sbCwdBg, priority: prioCwd}) + } + if s.contextPct >= 0 { + // The context readout is the highlighted amber block on the right. + segs = append(segs, segment{text: s.ctxText(), fg: sbCtxFg, bg: sbCtxBg, priority: prioToken}) + } + if s.task != "" { + segs = append(segs, segment{text: glyphTask + " " + s.task, fg: sbTaskFg, bg: sbTaskBg, priority: prioTask}) + } + return segs +} + +// gitText formats the git segment, e.g. "⎇ master ●3 ⇡4": branch, then "●N" for +// N dirty entries and "⇡N" for N commits ahead, each shown only when non-zero. +func (s statusBar) gitText() string { + var b strings.Builder + b.WriteString(glyphGit + " " + s.git.branch) + if s.git.dirty > 0 { + fmt.Fprintf(&b, " %s%d", glyphDirty, s.git.dirty) + } + if s.git.ahead > 0 { + fmt.Fprintf(&b, " %s%d", glyphAhead, s.git.ahead) + } + return b.String() +} + +// ctxText formats the context segment, e.g. "◔ 90,866 (46%)" when the token +// count is known, or "◔ 46%" before the first token count arrives. +func (s statusBar) ctxText() string { + if s.tokens > 0 { + return fmt.Sprintf("%s %s (%d%%)", glyphCtx, humanizeInt(s.tokens), s.contextPct) + } + return fmt.Sprintf("%s %d%%", glyphCtx, s.contextPct) +} + +// humanizeInt renders n with thousands separators, e.g. 90866 → "90,866". +func humanizeInt(n int) string { + s := fmt.Sprintf("%d", n) + neg := strings.HasPrefix(s, "-") + if neg { + s = s[1:] + } + var b strings.Builder + for i, r := range s { + if i > 0 && (len(s)-i)%3 == 0 { + b.WriteByte(',') + } + b.WriteRune(r) + } + if neg { + return "-" + b.String() + } + return b.String() +} + +// highestText returns the text of the highest-priority segment, used as the last +// thing standing when the terminal cannot even fit one full segment. +func (s statusBar) highestText() string { + segs := s.segments() + if len(segs) == 0 { + return "" + } + best := segs[0] + for _, seg := range segs[1:] { + if seg.priority > best.priority { + best = seg + } + } + return best.text +} + +// dropLowest removes one occurrence of the lowest-priority segment, preserving +// display order among the rest. It returns the shortened slice. +func dropLowest(segs []segment) []segment { + if len(segs) == 0 { + return segs + } + lowIdx := 0 + for i, seg := range segs { + if seg.priority < segs[lowIdx].priority { + lowIdx = i + } + } + out := make([]segment, 0, len(segs)-1) + out = append(out, segs[:lowIdx]...) + out = append(out, segs[lowIdx+1:]...) + return out +} + +// abbreviateHome replaces a leading $HOME in path with "~" so the status bar +// stays compact. It leaves paths outside $HOME untouched and never fails. +func abbreviateHome(path string) string { + home := homeDir() + if home == "" || path == "" { + return path + } + if path == home { + return "~" + } + if strings.HasPrefix(path, home+"/") { + return "~" + strings.TrimPrefix(path, home) + } + return path +} + +// homeDir returns the user's home directory, or "" when it cannot be +// determined. Kept as a tiny wrapper so abbreviateHome stays testable. +func homeDir() string { + h, err := os.UserHomeDir() + if err != nil { + return "" + } + return h +} diff --git a/pigo/internal/cli/tui/statusbar_test.go b/pigo/internal/cli/tui/statusbar_test.go new file mode 100644 index 0000000..5ec870b --- /dev/null +++ b/pigo/internal/cli/tui/statusbar_test.go @@ -0,0 +1,164 @@ +package tui + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/ui" +) + +// newTestStatusBar builds a status bar with a known cwd (already ~-abbreviated +// by the caller's intent) so tests do not depend on the real $HOME. +func newTestStatusBar() statusBar { + opts := Options{Model: "claude-opus", ThinkingLevel: agentcore.ThinkingHigh} + s := newStatusBar(DefaultTheme(), opts, "/tmp/project") + s.cwd = "~/project" + return s +} + +func TestStatusBarRendersAllFields(t *testing.T) { + s := newTestStatusBar() + s.SetGit(gitInfoMsg{branch: "master", dirty: 3, ahead: 4, ok: true}) + s.SetTelemetry(telemetryEventView{util: 0.42, window: 200000}) + s.SetTask("running: Read") + + const width = 200 + out := s.Render(width) + + for _, want := range []string{ + "pigo", // app badge + "claude-opus", // model + "high", // thinking level + "~/project", // cwd + "master", // git branch + glyphDirty + "3", // dirty marker + glyphAhead + "4", // ahead marker + "42%", // context usage + "running: Read", // task + } { + if !strings.Contains(out, want) { + t.Errorf("render missing %q; got %q", want, out) + } + } + + if w := ui.Width(out); w > width { + t.Errorf("render width %d exceeds terminal width %d", w, width) + } +} + +func TestStatusBarHidesGitWhenNotRepo(t *testing.T) { + s := newTestStatusBar() + s.SetGit(gitInfoMsg{ok: false}) + + out := s.Render(120) + if strings.Contains(out, "master") || strings.Contains(out, "*") || strings.Contains(out, "+") { + t.Errorf("git segment should be hidden when ok=false; got %q", out) + } +} + +func TestStatusBarHidesContextWhenUnknown(t *testing.T) { + s := newTestStatusBar() + // No telemetry set (window 0) → context segment hidden. + s.SetTelemetry(telemetryEventView{util: 0.5, window: 0}) + + out := s.Render(120) + if strings.Contains(out, glyphCtx) { + t.Errorf("context segment should be hidden when window unknown; got %q", out) + } +} + +func TestStatusBarTruncationKeepsPriorityFields(t *testing.T) { + s := newTestStatusBar() + s.SetGit(gitInfoMsg{branch: "master", dirty: 3, ahead: 4, ok: true}) + s.SetTelemetry(telemetryEventView{util: 0.42, window: 200000}) + s.SetTask("TASK") + + // Narrow width: only the highest-priority fields (task > model > token) + // should survive; cwd and git should drop first. + const width = 24 + out := s.Render(width) + + if w := ui.Width(out); w > width { + t.Fatalf("truncated render width %d exceeds %d: %q", w, width, out) + } + if !strings.Contains(out, "TASK") { + t.Errorf("highest-priority task field dropped under truncation: %q", out) + } + // cwd (lowest priority) must be gone before task. + if strings.Contains(out, "~/project") { + t.Errorf("lowest-priority cwd should drop first under truncation: %q", out) + } +} + +func TestStatusBarVeryNarrowNeverOverflows(t *testing.T) { + s := newTestStatusBar() + s.SetTask("a-fairly-long-task-description-that-cannot-fit") + + for _, width := range []int{1, 2, 3, 5, 8} { + out := s.Render(width) + if w := ui.Width(out); w > width { + t.Errorf("width %d: render width %d overflows: %q", width, w, out) + } + } +} + +func TestStatusBarZeroWidthEmpty(t *testing.T) { + s := newTestStatusBar() + if out := s.Render(0); out != "" { + t.Errorf("zero width should render empty, got %q", out) + } +} + +// TestStatusBarContextTokenCount checks the context segment shows a +// comma-grouped token count with the percentage once telemetry reports tokens. +func TestStatusBarContextTokenCount(t *testing.T) { + s := newTestStatusBar() + s.SetTelemetry(telemetryEventView{util: 0.46, window: 200000, tokens: 90866}) + + out := s.Render(200) + for _, want := range []string{glyphCtx, "90,866", "46%"} { + if !strings.Contains(out, want) { + t.Errorf("render missing %q; got %q", want, out) + } + } +} + +func TestHumanizeInt(t *testing.T) { + cases := map[int]string{0: "0", 90866: "90,866", 1000: "1,000", 999: "999", 1234567: "1,234,567"} + for in, want := range cases { + if got := humanizeInt(in); got != want { + t.Errorf("humanizeInt(%d) = %q, want %q", in, got, want) + } + } +} + +func TestAbbreviateHome(t *testing.T) { + home := homeDir() + if home == "" { + t.Skip("no home dir available") + } + if got := abbreviateHome(home); got != "~" { + t.Errorf("abbreviateHome(home) = %q, want ~", got) + } + if got := abbreviateHome(home + "/foo/bar"); got != "~/foo/bar" { + t.Errorf("abbreviateHome(home/foo/bar) = %q, want ~/foo/bar", got) + } + if got := abbreviateHome("/etc/passwd"); got != "/etc/passwd" { + t.Errorf("abbreviateHome(/etc/passwd) = %q, want unchanged", got) + } +} + +// TestStatusBarGitTextFormatting checks the "*N +N" markers appear only when +// non-zero. +func TestStatusBarGitTextFormatting(t *testing.T) { + s := newTestStatusBar() + s.SetGit(gitInfoMsg{branch: "main", ok: true}) + out := s.Render(120) + if strings.Contains(out, "*") || strings.Contains(out, "+") { + t.Errorf("clean tree should show no *N/+N markers: %q", out) + } + if !strings.Contains(out, "main") { + t.Errorf("branch name missing: %q", out) + } +} diff --git a/pigo/internal/cli/tui/subagentpanel.go b/pigo/internal/cli/tui/subagentpanel.go new file mode 100644 index 0000000..12ad0b9 --- /dev/null +++ b/pigo/internal/cli/tui/subagentpanel.go @@ -0,0 +1,317 @@ +package tui + +import ( + "fmt" + "strings" + "time" +) + +// This file renders the multi-line sub-agent status panel (SPEC 4.4, US-006): a +// block shown just above the working spinner while one or more sub-agents +// dispatched by the `task` tool are running. Each active sub-agent contributes +// exactly one status line of the form: +// +// ⏺ {desc} · {activity} ({elapsed} · ↓{tokens}) +// +// The panel is also interactive: while the input box is empty, ↑/↓ move a +// selection cursor over the rows and Enter expands the selected row to show that +// sub-agent's accumulated text output inline (below its status line), Esc +// collapses. The panel is a pure function of the model's ordered active-subagent +// set plus its selection state; it is re-rendered every spinner tick so the +// elapsed clock stays live without a dedicated timer. When there are no active +// sub-agents it renders nothing (zero lines, zero height), leaving the existing +// single-run layout untouched. + +// maxExpandedLines caps how many wrapped output lines an expanded row shows. The +// output can grow without bound, so only the most recent lines are kept visible; +// older content scrolls off the top of the inline pane. +const maxExpandedLines = 12 + +// subagentRow is one live sub-agent's status, keyed by the parent task tool-call +// id. start is recorded when the row is added so elapsed can be computed at +// render time; activity/tokens are refreshed by subagentProgressMsg; output +// accumulates the sub-agent's forwarded text (toolUpdate deltas + final result) +// for the inline expanded view. +type subagentRow struct { + id string + desc string + activity string + tokens int + start time.Time + output string +} + +// subagentPanel is the ordered set of live sub-agents. order preserves insertion +// order (so rows render stably, oldest first) while byID gives O(1) lookup for +// progress updates and removal. selecting reports whether a row is cursored; +// selected is that row's index into order (meaningful only while selecting is +// true); expanded reports whether the selected row shows its output inline. The +// zero value is a valid empty, unselected panel — selecting defaults false so the +// selected int's zero value never spuriously marks row 0. +type subagentPanel struct { + order []string + byID map[string]*subagentRow + selecting bool + selected int + expanded bool +} + +// add records a newly dispatched sub-agent (a toolStartMsg with name=="task"). +// It is idempotent on the id: a duplicate start refreshes the description and +// resets the start clock rather than adding a second row. +func (p *subagentPanel) add(id, desc string, now time.Time) { + if p.byID == nil { + p.byID = make(map[string]*subagentRow) + } + if row, ok := p.byID[id]; ok { + row.desc = desc + row.start = now + return + } + p.byID[id] = &subagentRow{id: id, desc: desc, start: now} + p.order = append(p.order, id) +} + +// update folds a progress event into the row for id, refreshing its activity and +// token estimate. A progress for an unknown id (late/out-of-order, arriving +// before or without a start) adds the row so no update is lost; now seeds its +// start clock in that case. +func (p *subagentPanel) update(id, desc, activity string, tokens int, now time.Time) { + if p.byID == nil { + p.byID = make(map[string]*subagentRow) + } + row, ok := p.byID[id] + if !ok { + row = &subagentRow{id: id, desc: desc, start: now} + p.byID[id] = row + p.order = append(p.order, id) + } + if activity != "" { + row.activity = activity + } + if desc != "" { + row.desc = desc + } + row.tokens = tokens +} + +// appendOutput accumulates a forwarded text delta into the row for id, so the +// expanded view can show the sub-agent's running output. Deltas for an unknown id +// are ignored (the row's start/end brackets its output; nothing to attach to). +func (p *subagentPanel) appendOutput(id, delta string) { + if delta == "" { + return + } + if row, ok := p.byID[id]; ok { + row.output += delta + } +} + +// remove drops the row for id (the task's toolEndMsg). It is a no-op when id is +// absent, so an end without a matching start — or a duplicate end — is safe. +// The selection is clamped to the shrunken order so the cursor never dangles past +// the end; removing the last row clears the selection entirely. +func (p *subagentPanel) remove(id string) { + if _, ok := p.byID[id]; !ok { + return + } + delete(p.byID, id) + for i, v := range p.order { + if v == id { + p.order = append(p.order[:i], p.order[i+1:]...) + break + } + } + if len(p.order) == 0 { + p.clearSelection() + return + } + if p.selected >= len(p.order) { + p.selected = len(p.order) - 1 + } +} + +// active reports the number of live sub-agents (status rows the panel would +// render), ignoring any extra rows an expanded row contributes. +func (p *subagentPanel) active() int { return len(p.order) } + +// hasSelection reports whether a row is currently cursored. +func (p *subagentPanel) hasSelection() bool { + return p.selecting && p.selected >= 0 && p.selected < len(p.order) +} + +// clearSelection drops the cursor and collapses any expansion. +func (p *subagentPanel) clearSelection() { + p.selecting = false + p.selected = 0 + p.expanded = false +} + +// selectUp moves the cursor to the previous row. With no current selection the +// first press lands on the last (bottom-most) row; moving up collapses any open +// expansion so it re-anchors to the newly selected row. +func (p *subagentPanel) selectUp() { + if len(p.order) == 0 { + return + } + if !p.selecting { + p.selecting = true + p.selected = len(p.order) - 1 + } else if p.selected > 0 { + p.selected-- + } + p.expanded = false +} + +// selectDown moves the cursor to the next row. With no current selection the +// first press lands on the first (top-most) row; moving down collapses any open +// expansion so it re-anchors to the newly selected row. +func (p *subagentPanel) selectDown() { + if len(p.order) == 0 { + return + } + if !p.selecting { + p.selecting = true + p.selected = 0 + } else if p.selected < len(p.order)-1 { + p.selected++ + } + p.expanded = false +} + +// toggleExpand flips the expanded state of the selected row. It is a no-op when +// nothing is selected. +func (p *subagentPanel) toggleExpand() { + if p.hasSelection() { + p.expanded = !p.expanded + } +} + +// expandedID returns the id of the currently expanded row, or "" when no row is +// expanded. It lets the model relayout only when a streamed delta lands on the +// row whose inline output pane is on screen. +func (p *subagentPanel) expandedID() string { + if p.expanded && p.hasSelection() { + return p.order[p.selected] + } + return "" +} + +// lineCount reports how many terminal rows the panel occupies at the given width: +// one status line per active sub-agent, plus the wrapped output lines when the +// selected row is expanded. relayout uses this to reserve exactly the right +// height so the transcript never overlaps the panel. +func (p subagentPanel) lineCount(width int) int { + if len(p.order) == 0 || width <= 0 { + return 0 + } + n := len(p.order) + if p.expanded && p.hasSelection() { + if row := p.byID[p.order[p.selected]]; row != nil { + n += len(p.expandedLines(row, width)) + } + } + return n +} + +// view renders the panel to a string, one status line per active sub-agent in +// insertion order, each truncated to width display columns. The selected row is +// marked with a leading cursor and, when expanded, its accumulated output is +// rendered on the following (indented, wrapped) lines. It returns "" when there +// are no active sub-agents or width is non-positive, so an empty panel +// contributes zero rows and zero height. now is the reference time elapsed is +// measured from (the spinner tick's time) so the clock advances each frame. +func (p subagentPanel) view(theme Theme, width int, now time.Time) string { + if len(p.order) == 0 || width <= 0 { + return "" + } + lines := make([]string, 0, len(p.order)+1) + for i, id := range p.order { + row := p.byID[id] + if row == nil { + continue + } + cursored := p.selecting && i == p.selected + lines = append(lines, TruncateToWidth(row.render(theme, now, cursored), width)) + if cursored && p.expanded { + for _, out := range p.expandedLines(row, width) { + lines = append(lines, theme.System.Render(out)) + } + } + } + return strings.Join(lines, "\n") +} + +// expandedLines builds the wrapped, indented, tail-capped output lines shown +// under an expanded row. An empty output yields a single placeholder line so the +// pane is never blank. Lines are already truncated to width (as plain text); the +// caller styles them. +func (p subagentPanel) expandedLines(row *subagentRow, width int) []string { + const indent = " " + out := strings.TrimRight(row.output, "\n") + if out == "" { + return []string{indent + "(no output yet)"} + } + var wrapped []string + for _, para := range strings.Split(out, "\n") { + for _, seg := range wrapToWidth(para, width-len(indent)) { + wrapped = append(wrapped, indent+seg) + } + } + if len(wrapped) > maxExpandedLines { + wrapped = wrapped[len(wrapped)-maxExpandedLines:] + } + return wrapped +} + +// wrapToWidth breaks s into segments no wider than width display columns, cutting +// on the column boundary (there is no word-aware wrapping here — sub-agent output +// is arbitrary text/code). An empty line yields one empty segment so blank lines +// in the output are preserved. +func wrapToWidth(s string, width int) []string { + if width <= 0 { + return []string{s} + } + if s == "" { + return []string{""} + } + var segs []string + for s != "" { + seg := TruncateToWidth(s, width) + if seg == "" { // guard against no forward progress on odd-width runes + segs = append(segs, s) + break + } + segs = append(segs, seg) + s = s[len(seg):] + } + return segs +} + +// render builds one status line for a row: "{cursor}⏺ {desc} · {activity} +// ({elapsed} · ↓{tokens})". A blank description is omitted (the line leads with +// the glyph and activity); a zero token estimate drops the "↓" stat. When +// selected, the line leads with a "❯ " cursor and the head takes the accent color +// to stand out; otherwise the glyph + head take the spinner color and the +// parenthetical stats are dim, mirroring the spinner line. +func (r subagentRow) render(theme Theme, now time.Time, selected bool) string { + var head strings.Builder + head.WriteString("⏺") + if r.desc != "" { + fmt.Fprintf(&head, " %s ·", r.desc) + } + fmt.Fprintf(&head, " %s", r.activity) + + stats := formatElapsed(now.Sub(r.start)) + if r.tokens > 0 { + stats += " · ↓" + humanizeInt(r.tokens) + } + + headStyle := theme.Spinner + cursor := " " + if selected { + headStyle = theme.Accent + cursor = theme.Accent.Render("❯ ") + } + return cursor + headStyle.Render(head.String()) + " " + theme.System.Render("("+stats+")") +} diff --git a/pigo/internal/cli/tui/subagentpanel_test.go b/pigo/internal/cli/tui/subagentpanel_test.go new file mode 100644 index 0000000..1543eae --- /dev/null +++ b/pigo/internal/cli/tui/subagentpanel_test.go @@ -0,0 +1,292 @@ +package tui + +import ( + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// TestSubagentPanelLifecycle exercises the ordered add/update/remove set: rows +// keep insertion order, a progress refreshes activity/tokens, a progress for an +// unknown id adds a row (late/out-of-order safe), and remove drops exactly one +// row (a no-op for an absent id). +func TestSubagentPanelLifecycle(t *testing.T) { + now := time.Now() + var p subagentPanel + + p.add("a", "task A", now) + p.add("b", "task B", now) + if got := p.active(); got != 2 { + t.Fatalf("active after two adds = %d, want 2", got) + } + if p.order[0] != "a" || p.order[1] != "b" { + t.Errorf("order = %v, want [a b]", p.order) + } + + p.update("a", "task A", "Editing", 120, now) + if row := p.byID["a"]; row.activity != "Editing" || row.tokens != 120 { + t.Errorf("row a after update = %+v, want activity=Editing tokens=120", row) + } + + // A progress for an id that never started adds the row (SPEC 5.4). + p.update("c", "task C", "Reading", 0, now) + if got := p.active(); got != 3 { + t.Fatalf("active after late progress = %d, want 3", got) + } + if p.order[2] != "c" { + t.Errorf("order = %v, want c appended last", p.order) + } + + // Removing a middle row preserves the order of the rest. + p.remove("a") + if got := p.active(); got != 2 { + t.Fatalf("active after remove = %d, want 2", got) + } + if p.order[0] != "b" || p.order[1] != "c" { + t.Errorf("order after remove(a) = %v, want [b c]", p.order) + } + + // Removing an absent id is a no-op. + p.remove("zzz") + if got := p.active(); got != 2 { + t.Errorf("active after remove(absent) = %d, want 2", got) + } +} + +// TestSubagentPanelEmptyView verifies an empty panel renders nothing — zero +// lines, zero height — so the single-run layout is untouched. +func TestSubagentPanelEmptyView(t *testing.T) { + var p subagentPanel + if got := p.view(DefaultTheme(), 80, time.Now()); got != "" { + t.Errorf("empty panel view = %q, want empty", got) + } + // A non-empty panel with a non-positive width also renders nothing. + p.add("a", "task", time.Now()) + if got := p.view(DefaultTheme(), 0, time.Now()); got != "" { + t.Errorf("zero-width view = %q, want empty", got) + } +} + +// TestSubagentPanelViewLines verifies one line per active sub-agent, each +// carrying the description, activity, and token stat. +func TestSubagentPanelViewLines(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", "build parser", now) + p.update("a", "build parser", "Editing", 1200, now) + p.add("b", "run tests", now) + p.update("b", "run tests", "Running bash", 0, now) + + view := p.view(DefaultTheme(), 200, now.Add(65*time.Second)) + lines := strings.Split(view, "\n") + if len(lines) != 2 { + t.Fatalf("view has %d lines, want 2: %q", len(lines), view) + } + if !strings.Contains(lines[0], "build parser") || !strings.Contains(lines[0], "Editing") { + t.Errorf("line[0] = %q, want desc + activity", lines[0]) + } + if !strings.Contains(lines[0], "1m 5s") { + t.Errorf("line[0] = %q, want elapsed 1m 5s", lines[0]) + } + if !strings.Contains(lines[0], "1,200") { + t.Errorf("line[0] = %q, want token stat 1,200", lines[0]) + } + if !strings.Contains(lines[1], "run tests") || !strings.Contains(lines[1], "Running bash") { + t.Errorf("line[1] = %q, want desc + activity", lines[1]) + } + // A zero token estimate omits the ↓ stat. + if strings.Contains(lines[1], "↓") { + t.Errorf("line[1] = %q, should omit ↓ for zero tokens", lines[1]) + } +} + +// TestSubagentPanelViewTruncation verifies each rendered line is clipped to the +// given terminal width (display columns), never exceeding it. +func TestSubagentPanelViewTruncation(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", strings.Repeat("very long description ", 10), now) + p.update("a", "", "Searching", 42, now) + + const width = 30 + view := p.view(DefaultTheme(), width, now) + for _, line := range strings.Split(view, "\n") { + if w := ui.Width(line); w > width { + t.Errorf("line width = %d, want <= %d: %q", w, width, line) + } + } +} + +// TestSubagentPanelViewBlankDescription verifies a row with no description still +// renders (leading with the activity) rather than producing a dangling "·". +func TestSubagentPanelViewBlankDescription(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", "", now) + p.update("a", "", "Thinking", 0, now) + line := p.view(DefaultTheme(), 100, now) + if !strings.Contains(line, "Thinking") { + t.Errorf("view = %q, want activity", line) + } + if strings.Contains(line, " · Thinking") { + t.Errorf("view = %q, blank desc should not leave a leading ' · '", line) + } +} + +// TestSubagentPanelSelection verifies the cursor navigation: a fresh panel has no +// selection, the first ↓ lands on the top row and the first ↑ on the bottom row, +// movement clamps at both ends, and clearSelection resets to no-cursor state. +func TestSubagentPanelSelection(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", "task A", now) + p.add("b", "task B", now) + p.add("c", "task C", now) + + if p.hasSelection() { + t.Fatal("fresh panel should have no selection") + } + + // First ↓ selects the top row; further ↓ advance and clamp at the bottom. + p.selectDown() + if !p.hasSelection() || p.selected != 0 { + t.Fatalf("after first down: hasSelection=%v selected=%d, want true/0", p.hasSelection(), p.selected) + } + p.selectDown() + p.selectDown() + p.selectDown() // clamp + if p.selected != 2 { + t.Errorf("selected after clamp down = %d, want 2", p.selected) + } + + // ↑ retreats and clamps at the top. + p.selectUp() + if p.selected != 1 { + t.Errorf("selected after up = %d, want 1", p.selected) + } + p.selectUp() + p.selectUp() // clamp + if p.selected != 0 { + t.Errorf("selected after clamp up = %d, want 0", p.selected) + } + + p.clearSelection() + if p.hasSelection() { + t.Error("clearSelection should drop the cursor") + } + + // From no selection, the first ↑ lands on the bottom row. + p.selectUp() + if !p.hasSelection() || p.selected != 2 { + t.Errorf("first up from none: selected=%d, want 2", p.selected) + } +} + +// TestSubagentPanelExpandView verifies that expanding a selected row appends its +// accumulated output below the status line, that the cursor marker is present, +// that lineCount matches the rendered height, and that collapsing removes the +// extra lines. +func TestSubagentPanelExpandView(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", "task A", now) + p.add("b", "task B", now) + p.appendOutput("b", "hello from B\nsecond line") + + p.selectDown() // selects "a" + p.selectDown() // selects "b" + if id := p.expandedID(); id != "" { + t.Errorf("expandedID before toggle = %q, want empty", id) + } + p.toggleExpand() + if id := p.expandedID(); id != "b" { + t.Errorf("expandedID after toggle = %q, want b", id) + } + + const width = 80 + view := p.view(DefaultTheme(), width, now) + lines := strings.Split(view, "\n") + if got := p.lineCount(width); got != len(lines) { + t.Errorf("lineCount = %d, rendered %d lines", got, len(lines)) + } + // Two status rows + two output lines. + if len(lines) != 4 { + t.Fatalf("expanded view has %d lines, want 4: %q", len(lines), view) + } + if !strings.Contains(view, "❯") { + t.Errorf("expanded view missing selection cursor: %q", view) + } + if !strings.Contains(view, "hello from B") || !strings.Contains(view, "second line") { + t.Errorf("expanded view missing output: %q", view) + } + + // Collapsing reclaims the output lines. + p.toggleExpand() + if got := p.lineCount(width); got != 2 { + t.Errorf("lineCount after collapse = %d, want 2", got) + } +} + +// TestSubagentPanelExpandTruncates verifies the inline output is wrapped and +// tail-capped: every rendered line fits the width and no more than +// maxExpandedLines output lines are shown. +func TestSubagentPanelExpandTruncates(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", "task A", now) + // Many short lines exceed the cap; one very long line must wrap. + p.appendOutput("a", strings.Repeat("x\n", 40)) + p.appendOutput("a", strings.Repeat("y", 500)) + p.selectDown() + p.toggleExpand() + + const width = 40 + view := p.view(DefaultTheme(), width, now) + lines := strings.Split(view, "\n") + for _, line := range lines { + if w := ui.Width(line); w > width { + t.Errorf("line width = %d, want <= %d: %q", w, width, line) + } + } + // 1 status row + at most maxExpandedLines output lines. + if len(lines) > 1+maxExpandedLines { + t.Errorf("expanded view has %d lines, want <= %d", len(lines), 1+maxExpandedLines) + } +} + +// TestSubagentPanelRemoveClampsSelection verifies that removing rows keeps the +// selection valid: removing the selected bottom row clamps the cursor to the new +// last row, and removing the final row clears the selection entirely. +func TestSubagentPanelRemoveClampsSelection(t *testing.T) { + now := time.Now() + var p subagentPanel + p.add("a", "task A", now) + p.add("b", "task B", now) + p.selectDown() + p.selectDown() // selects "b" (index 1) + + p.remove("b") + if !p.hasSelection() || p.selected != 0 { + t.Errorf("after remove(b): hasSelection=%v selected=%d, want true/0", p.hasSelection(), p.selected) + } + + p.remove("a") + if p.hasSelection() { + t.Error("removing the last row should clear the selection") + } + if got := p.lineCount(80); got != 0 { + t.Errorf("empty panel lineCount = %d, want 0", got) + } +} + +// TestSubagentPanelAppendOutputUnknown verifies deltas for an unknown id are +// dropped (no phantom row, no panic). +func TestSubagentPanelAppendOutputUnknown(t *testing.T) { + var p subagentPanel + p.appendOutput("ghost", "data") // must not panic or add a row + if p.active() != 0 { + t.Errorf("active after appendOutput to unknown id = %d, want 0", p.active()) + } +} diff --git a/pigo/internal/cli/tui/theme.go b/pigo/internal/cli/tui/theme.go new file mode 100644 index 0000000..9f7ca90 --- /dev/null +++ b/pigo/internal/cli/tui/theme.go @@ -0,0 +1,195 @@ +package tui + +import ( + "strings" + + "charm.land/lipgloss/v2" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// Theme bundles the lipgloss styles for every visual element the TUI paints so +// the transcript, tool cards and status bar share one palette instead of each +// call site hand-rolling colors (see tasks/spec-tui-agent.md Sections 2.2, 5.1). +// The reference palette is: success green, error/warn red & yellow, file/accent +// blue, and gray for secondary chrome. Styles are plain value types, so a Theme +// is cheap to copy and safe to pass by value. +type Theme struct { + // User styles the human's turns in the transcript. + User lipgloss.Style + // Assistant styles the model's turns in the transcript. + Assistant lipgloss.Style + // System styles system / meta notices (secondary gray). + System lipgloss.Style + // ToolHeader styles the title line of a tool invocation card. + ToolHeader lipgloss.Style + // ToolBody styles the body/output region of a tool card. + ToolBody lipgloss.Style + // StatusBar styles the persistent bottom status bar. + StatusBar lipgloss.Style + // Accent styles file names and other highlighted tokens (blue). + Accent lipgloss.Style + // Error styles failure messages (red). + Error lipgloss.Style + // Warn styles warnings (yellow). + Warn lipgloss.Style + // Success styles successful outcomes (green). + Success lipgloss.Style + // ScrollThumb styles the transcript scrollbar thumb (medium gray block). + ScrollThumb lipgloss.Style + // ScrollTrack styles the transcript scrollbar track (dim shaded column). + ScrollTrack lipgloss.Style + // Spinner styles the animated "working" indicator glyph + verb (warm coral). + Spinner lipgloss.Style +} + +// Palette color numbers use the ANSI 256-color cube so the theme renders +// consistently across terminals without depending on true-color support. +const ( + colorSuccess = "42" // green + colorError = "196" // red + colorWarn = "214" // yellow/amber + colorAccent = "39" // blue (file names, highlights) + colorGray = "245" // secondary / muted text + colorScroll = "250" // scrollbar thumb (bright gray pill, clearly visible) + colorTrack = "240" // scrollbar groove (dim gray, visible but recessive) + colorUser = "15" // bright white + colorAssist = "252" // near-white + colorStatus = "62" // status bar background (violet) + colorSpinner = "173" // spinner glyph/verb (warm coral, matches Claude Code) +) + +// DefaultTheme returns the built-in palette described in the SPEC: success +// green, error/warn red & yellow, file/accent blue, and gray for secondary +// chrome. It performs no I/O and never panics, so callers can construct it +// eagerly at startup. +func DefaultTheme() Theme { + return Theme{ + User: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorUser)). + Bold(true), + Assistant: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorAssist)), + System: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorGray)). + Italic(true), + ToolHeader: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorAccent)). + Bold(true), + ToolBody: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorGray)), + StatusBar: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorUser)). + Background(lipgloss.Color(colorStatus)), + Accent: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorAccent)), + Error: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorError)). + Bold(true), + Warn: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorWarn)), + Success: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorSuccess)), + ScrollThumb: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorScroll)), + ScrollTrack: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorTrack)), + Spinner: lipgloss.NewStyle(). + Foreground(lipgloss.Color(colorSpinner)). + Bold(true), + } +} + +// ellipsis is the single rune appended to truncated strings. It is itself a +// single display column, so it is cheap to reserve room for. +const ellipsis = "…" + +// WrapToWidth wraps s to at most width display columns per line, measuring width +// by terminal cells (CJK and emoji count as two) via ui.Width rather than byte +// length. It never splits inside a multi-byte rune or a double-width character: +// a rune that would overflow the current line starts a new line instead. Any +// existing newlines in s are preserved as hard breaks. A non-positive width is +// treated as "no wrapping" and s is returned unchanged. +func WrapToWidth(s string, width int) string { + if width <= 0 { + return s + } + + var out strings.Builder + lines := strings.Split(s, "\n") + for li, line := range lines { + if li > 0 { + out.WriteByte('\n') + } + wrapLine(&out, line, width) + } + return out.String() +} + +// wrapLine wraps a single newline-free line into out, breaking on rune +// boundaries so no double-width rune is ever cut in half. +func wrapLine(out *strings.Builder, line string, width int) { + cur := 0 // display width accumulated on the current output line + first := true + for _, r := range line { + rw := ui.Width(string(r)) + if !first && cur+rw > width { + out.WriteByte('\n') + cur = 0 + } + out.WriteRune(r) + cur += rw + first = false + } +} + +// TruncateToWidth returns s clipped to at most width display columns, appending +// an ellipsis "…" when it removes content. Width is measured in terminal cells +// (CJK and emoji count as two) via ui.Width, and truncation happens on rune +// boundaries so a double-width character is never sliced. The returned string's +// display width is guaranteed to be <= width. A non-positive width yields the +// empty string. +func TruncateToWidth(s string, width int) string { + if width <= 0 { + return "" + } + if ui.Width(s) <= width { + return s + } + + // Reserve room for the ellipsis. If width is too small to even hold the + // ellipsis plus one column, fall back to fitting bare runes into width. + budget := width - ui.Width(ellipsis) + if budget <= 0 { + return fitRunes(s, width) + } + + var b strings.Builder + used := 0 + for _, r := range s { + rw := ui.Width(string(r)) + if used+rw > budget { + break + } + b.WriteRune(r) + used += rw + } + b.WriteString(ellipsis) + return b.String() +} + +// fitRunes packs as many leading runes of s as fit within width columns without +// any ellipsis, breaking on rune boundaries. +func fitRunes(s string, width int) string { + var b strings.Builder + used := 0 + for _, r := range s { + rw := ui.Width(string(r)) + if used+rw > width { + break + } + b.WriteRune(r) + used += rw + } + return b.String() +} diff --git a/pigo/internal/cli/tui/theme_test.go b/pigo/internal/cli/tui/theme_test.go new file mode 100644 index 0000000..027c565 --- /dev/null +++ b/pigo/internal/cli/tui/theme_test.go @@ -0,0 +1,116 @@ +package tui + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +func TestDefaultThemeRenders(t *testing.T) { + th := DefaultTheme() + + cases := map[string]string{ + "user": th.User.Render("hi"), + "assistant": th.Assistant.Render("ok"), + "system": th.System.Render("note"), + "toolHeader": th.ToolHeader.Render("Bash"), + "toolBody": th.ToolBody.Render("output"), + "statusBar": th.StatusBar.Render("status"), + "accent": th.Accent.Render("file.go"), + "error": th.Error.Render("boom"), + "warn": th.Warn.Render("careful"), + "success": th.Success.Render("done"), + } + for name, got := range cases { + if got == "" { + t.Errorf("style %s rendered empty output", name) + } + } +} + +func TestWrapToWidthDisplayWidth(t *testing.T) { + // Mix double-width CJK, an emoji, and ASCII. + const input = "你好world世界🚀测试abc" + const width = 6 + + wrapped := WrapToWidth(input, width) + + // Reassembling the wrapped lines (minus the inserted newlines) must equal + // the original: nothing is dropped or split inside a rune. + if got := strings.ReplaceAll(wrapped, "\n", ""); got != input { + t.Fatalf("wrap altered content: got %q want %q", got, input) + } + + for _, line := range strings.Split(wrapped, "\n") { + if w := ui.Width(line); w > width { + t.Errorf("line %q has display width %d > %d", line, w, width) + } + // Guard against a mid-rune cut producing invalid UTF-8. + if !isValidBoundary(line) { + t.Errorf("line %q was cut inside a multibyte rune", line) + } + } +} + +func TestWrapToWidthPreservesNewlines(t *testing.T) { + out := WrapToWidth("ab\ncd", 10) + if out != "ab\ncd" { + t.Fatalf("wrap collapsed existing newlines: got %q", out) + } +} + +func TestWrapToWidthNonPositive(t *testing.T) { + const s = "你好world" + if got := WrapToWidth(s, 0); got != s { + t.Errorf("width<=0 should return input unchanged, got %q", got) + } +} + +func TestTruncateToWidth(t *testing.T) { + tests := []struct { + name string + in string + width int + }{ + {"cjk", "你好世界测试内容很长", 6}, + {"emoji", "🚀🚀🚀🚀🚀🚀", 5}, + {"mixed", "abc你好def世界🚀tail", 8}, + {"ascii", "helloworld", 4}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := TruncateToWidth(tt.in, tt.width) + if w := ui.Width(got); w > tt.width { + t.Errorf("truncated %q to display width %d > %d (result %q)", tt.in, w, tt.width, got) + } + if !isValidBoundary(got) { + t.Errorf("truncation cut inside a multibyte rune: %q", got) + } + // It must actually be a truncation: contain the ellipsis when the + // input was wider than the budget. + if ui.Width(tt.in) > tt.width && !strings.Contains(got, ellipsis) { + t.Errorf("expected ellipsis in truncated result, got %q", got) + } + }) + } +} + +func TestTruncateToWidthNoTruncationNeeded(t *testing.T) { + const s = "你好" + if got := TruncateToWidth(s, 10); got != s { + t.Errorf("short string should be returned unchanged, got %q", got) + } +} + +func TestTruncateToWidthNonPositive(t *testing.T) { + if got := TruncateToWidth("你好", 0); got != "" { + t.Errorf("width<=0 should return empty string, got %q", got) + } +} + +// isValidBoundary reports whether s contains no invalid UTF-8, which would be +// the tell-tale of a cut inside a multibyte rune. +func isValidBoundary(s string) bool { + return strings.ToValidUTF8(s, "�") == s +} diff --git a/pigo/internal/cli/tui/toolcard.go b/pigo/internal/cli/tui/toolcard.go new file mode 100644 index 0000000..8113324 --- /dev/null +++ b/pigo/internal/cli/tui/toolcard.go @@ -0,0 +1,195 @@ +package tui + +import ( + "fmt" + "sort" + "strings" + + "charm.land/lipgloss/v2" + + "github.com/smallnest/pigo/internal/cli/ui" +) + +// This file implements the rich tool-call card component (US-006, SPEC 3.2, +// FR-6/7/8). A toolCard is a bordered inline block in the transcript that shows +// a single tool invocation: a header with the tool name and a status icon +// (running / success / warn), the decoded call arguments, and the tool's +// response rendered as an indented tree. Cards are created on toolStartMsg, +// completed on toolEndMsg, and toggled between a capped and a full response view +// with Ctrl+O (see model.go). All width math goes through ui.Width / +// WrapToWidth / TruncateToWidth so CJK and emoji (two columns) never split. + +// cardState is the lifecycle of a tool card: running while the tool executes, +// then success or warn once it finishes (warn covers a reported tool error). +type cardState int + +const ( + cardRunning cardState = iota + cardSuccess + cardWarn +) + +// respNode is one line of a tool's response, with depth giving the tree indent +// level (each level is rendered as two leading spaces). +type respNode struct { + text string + depth int +} + +// toolCard is a single tool invocation rendered as a bordered card. input holds +// the decoded call arguments (nil when the args were not a JSON object); +// response is the parsed result tree, populated on completion. expanded flips +// the response between a capped preview and the full tree. +type toolCard struct { + id string + name string + input map[string]any + response []respNode + state cardState + expanded bool +} + +// collapsedResponseLines is how many response lines a card shows before it is +// expanded; past this the preview is truncated and a Ctrl+O hint is appended. +const collapsedResponseLines = 5 + +// statusIcon returns the header status glyph for the card's state. Running is a +// spinner-like ellipsis, success a check, warn a bang. +func (c toolCard) statusIcon() string { + switch c.state { + case cardSuccess: + return "✓" + case cardWarn: + return "!" + default: + return "…" + } +} + +// styledIcon renders the status glyph with the state's theme color: gray while +// running, green on success, yellow/red on warn. +func (c toolCard) styledIcon(theme Theme) string { + icon := c.statusIcon() + switch c.state { + case cardSuccess: + return theme.Success.Render(icon) + case cardWarn: + return theme.Warn.Render(icon) + default: + return theme.System.Render(icon) + } +} + +// render draws the card at the given content width: a rounded border wrapping a +// header (status icon + tool name), an "Input arguments" section listing the input map, +// and a "Response" section with the tree lines. When not expanded the response +// is capped to collapsedResponseLines with a "(Ctrl+O for more)" hint; when +// expanded every line is shown. +func (c toolCard) render(theme Theme, width int) string { + if width < 4 { + width = 4 + } + // The rounded border consumes one column on each side; wrap everything to the + // inner width so nothing overflows the frame. + inner := width - 2 + + var lines []string + + icon := c.styledIcon(theme) + nameBudget := inner - ui.Width(icon) - 1 + if nameBudget < 1 { + nameBudget = 1 + } + header := c.name + if arg := c.primaryArg(); arg != "" { + header = c.name + "(" + arg + ")" + } + header = TruncateToWidth(header, nameBudget) + lines = append(lines, icon+" "+theme.ToolHeader.Render(header)) + + if len(c.input) > 0 { + lines = append(lines, theme.ToolBody.Render("Input arguments")) + for _, k := range sortedKeys(c.input) { + kv := " " + k + ": " + fmt.Sprintf("%v", c.input[k]) + lines = append(lines, theme.ToolBody.Render(WrapToWidth(kv, inner))) + } + } + + if len(c.response) > 0 { + lines = append(lines, theme.ToolBody.Render("Response")) + resp := c.response + truncated := false + if !c.expanded && len(resp) > collapsedResponseLines { + resp = resp[:collapsedResponseLines] + truncated = true + } + for _, n := range resp { + indent := strings.Repeat(" ", n.depth) + lines = append(lines, theme.ToolBody.Render(WrapToWidth(indent+n.text, inner))) + } + if truncated { + lines = append(lines, theme.System.Render("(Ctrl+O for more)")) + } + } + + border := lipgloss.NewStyle(). + Border(lipgloss.RoundedBorder()). + BorderForeground(lipgloss.Color(colorGray)). + Width(inner) + return border.Render(strings.Join(lines, "\n")) +} + +// primaryArg returns the most salient call argument to inline in the card header +// so the user can see what the tool is operating on at a glance (FR-6), e.g. +// Bash(cd /x && git add -A). It picks the command for bash and the file path for +// the file tools, otherwise the first argument in sorted-key order. Returns "" +// when the call carried no arguments. +func (c toolCard) primaryArg() string { + if len(c.input) == 0 { + return "" + } + var keyPrefs []string + switch strings.ToLower(c.name) { + case "bash": + keyPrefs = []string{"command"} + case "read", "write", "edit", "multiedit": + // The file tools emit "path"; accept "file_path" as a fallback for + // callers that use the Claude-style key. + keyPrefs = []string{"path", "file_path"} + } + for _, key := range keyPrefs { + if v, ok := c.input[key]; ok { + return fmt.Sprintf("%v", v) + } + } + keys := sortedKeys(c.input) + return fmt.Sprintf("%v", c.input[keys[0]]) +} + +// sortedKeys returns the map keys in a stable (sorted) order so the input +// section renders deterministically instead of in Go's random map order. +func sortedKeys(m map[string]any) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +// parseToolResult splits a tool's textual result into response tree nodes, +// inferring depth from leading whitespace (every two leading spaces is one +// level). Trailing empty lines are trimmed so the card does not render blank +// tail rows. +func parseToolResult(result string) []respNode { + lines := strings.Split(result, "\n") + for len(lines) > 0 && strings.TrimSpace(lines[len(lines)-1]) == "" { + lines = lines[:len(lines)-1] + } + nodes := make([]respNode, 0, len(lines)) + for _, ln := range lines { + leading := len(ln) - len(strings.TrimLeft(ln, " ")) + nodes = append(nodes, respNode{text: ln[leading:], depth: leading / 2}) + } + return nodes +} diff --git a/pigo/internal/cli/tui/toolcard_test.go b/pigo/internal/cli/tui/toolcard_test.go new file mode 100644 index 0000000..a4e9fe6 --- /dev/null +++ b/pigo/internal/cli/tui/toolcard_test.go @@ -0,0 +1,164 @@ +package tui + +import ( + "strings" + "testing" + + tea "charm.land/bubbletea/v2" +) + +// ctrlKey builds a Ctrl+ key press matching String()=="ctrl+". +func ctrlKey(r rune) tea.KeyPressMsg { + return tea.KeyPressMsg{Code: r, Mod: tea.ModCtrl} +} + +// TestParseToolResult verifies depth inference from leading spaces and trailing +// blank-line trimming. +func TestParseToolResult(t *testing.T) { + nodes := parseToolResult("root\n child\n grandchild\n\n") + if len(nodes) != 3 { + t.Fatalf("node count = %d, want 3 (trailing blank trimmed)", len(nodes)) + } + want := []respNode{ + {text: "root", depth: 0}, + {text: "child", depth: 1}, + {text: "grandchild", depth: 2}, + } + for i, w := range want { + if nodes[i] != w { + t.Errorf("node[%d] = %+v, want %+v", i, nodes[i], w) + } + } +} + +// TestToolCardRender checks the header (name + status icon), the input section, +// and the response tree lines appear in the rendered card, and that the status +// icon reflects the state. +func TestToolCardRender(t *testing.T) { + theme := DefaultTheme() + cases := []struct { + name string + state cardState + icon string + }{ + {"running", cardRunning, "…"}, + {"success", cardSuccess, "✓"}, + {"warn", cardWarn, "!"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + card := toolCard{ + id: "1", + name: "read_file", + input: map[string]any{"path": "/tmp/x"}, + response: parseToolResult("line one\n nested"), + state: tc.state, + } + out := card.render(theme, 60) + for _, want := range []string{"read_file", tc.icon, "Input arguments", "path: /tmp/x", "Response", "line one", "nested"} { + if !strings.Contains(out, want) { + t.Errorf("render missing %q\n%s", want, out) + } + } + }) + } +} + +// TestToolCardExpandTruncation verifies the collapsed card caps the response and +// shows the Ctrl+O hint, while the expanded card reveals every line. +func TestToolCardExpandTruncation(t *testing.T) { + theme := DefaultTheme() + var b strings.Builder + for i := 0; i < collapsedResponseLines+3; i++ { + b.WriteString("resp-line-") + b.WriteByte(byte('a' + i)) + b.WriteByte('\n') + } + card := toolCard{name: "grep", response: parseToolResult(b.String()), state: cardSuccess} + + collapsed := card.render(theme, 60) + if !strings.Contains(collapsed, "(Ctrl+O for more)") { + t.Errorf("collapsed card should show Ctrl+O hint\n%s", collapsed) + } + lastLine := "resp-line-" + string(byte('a'+collapsedResponseLines+2)) + if strings.Contains(collapsed, lastLine) { + t.Errorf("collapsed card should not show %q\n%s", lastLine, collapsed) + } + + card.expanded = true + expanded := card.render(theme, 60) + if strings.Contains(expanded, "(Ctrl+O for more)") { + t.Errorf("expanded card should not show Ctrl+O hint\n%s", expanded) + } + if !strings.Contains(expanded, lastLine) { + t.Errorf("expanded card should show %q\n%s", lastLine, expanded) + } +} + +// TestModelToolCardFlow drives the model through a tool start/end and asserts the +// card is created, transitions running→success, and that a failed tool yields +// warn. +func TestModelToolCardFlow(t *testing.T) { + m := NewModel(Options{}) + next, _ := m.Update(toolStartMsg{id: "t1", name: "read_file", input: map[string]any{"path": "a.go"}}) + mm := next.(Model) + card, ok := mm.toolCards["t1"] + if !ok { + t.Fatalf("toolStartMsg should create a card") + } + if card.state != cardRunning { + t.Errorf("new card state = %v, want cardRunning", card.state) + } + + next, _ = mm.Update(toolEndMsg{id: "t1", ok: true, result: "done\n detail"}) + mm = next.(Model) + if mm.toolCards["t1"].state != cardSuccess { + t.Errorf("state after ok end = %v, want cardSuccess", mm.toolCards["t1"].state) + } + if len(mm.toolCards["t1"].response) != 2 { + t.Errorf("response nodes = %d, want 2", len(mm.toolCards["t1"].response)) + } + + // A failed tool flips the same card to warn. + next, _ = m.Update(toolStartMsg{id: "t2", name: "bash"}) + mm = next.(Model) + next, _ = mm.Update(toolEndMsg{id: "t2", ok: false, result: "boom"}) + mm = next.(Model) + if mm.toolCards["t2"].state != cardWarn { + t.Errorf("state after failed end = %v, want cardWarn", mm.toolCards["t2"].state) + } +} + +// TestModelCtrlOTogglesExpanded verifies Ctrl+O flips the most-recent card's +// expanded flag so more response lines become visible. +func TestModelCtrlOTogglesExpanded(t *testing.T) { + m := NewModel(Options{}) + next, _ := m.Update(tea.WindowSizeMsg{Width: 60, Height: 24}) + mm := next.(Model) + + var b strings.Builder + for i := 0; i < collapsedResponseLines+3; i++ { + b.WriteString("row") + b.WriteByte(byte('0' + i)) + b.WriteByte('\n') + } + next, _ = mm.Update(toolStartMsg{id: "t1", name: "grep"}) + mm = next.(Model) + next, _ = mm.Update(toolEndMsg{id: "t1", ok: true, result: b.String()}) + mm = next.(Model) + + if mm.lastToolCard.expanded { + t.Fatalf("card should start collapsed") + } + next, _ = mm.Update(ctrlKey('o')) + mm = next.(Model) + if !mm.lastToolCard.expanded { + t.Errorf("Ctrl+O should expand the most-recent card") + } + // Toggling again collapses it. + next, _ = mm.Update(ctrlKey('o')) + mm = next.(Model) + if mm.lastToolCard.expanded { + t.Errorf("second Ctrl+O should collapse the card") + } +} diff --git a/pigo/internal/cli/tui/transcript.go b/pigo/internal/cli/tui/transcript.go new file mode 100644 index 0000000..a70fb2b --- /dev/null +++ b/pigo/internal/cli/tui/transcript.go @@ -0,0 +1,404 @@ +package tui + +import ( + "strings" + + "charm.land/bubbles/v2/viewport" + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// This file implements the scrolling transcript region of the full-screen TUI +// (US-005, SPEC 5.1 transcript, FR-5/FR-10). The transcript owns a +// viewport.Model and an ordered list of rendered blocks (user / assistant / +// system turns). Streaming assistant text arrives as textDeltaMsg values that +// append to the current assistant block; turnEndMsg finalizes it. Content is +// re-flowed through the viewport with theme.WrapToWidth at the live width so CJK +// and emoji never split mid-rune. Tool cards are a later node (#389); this file +// leaves a clean seam (system lines) without building cards. + +// blockRole distinguishes the three transcript block kinds so each renders with +// its own theme style. +type blockRole int + +const ( + roleUser blockRole = iota + roleAssistant + roleSystem + roleTool + // roleBanner is the startup logo + config splash. Its text is pre-rendered + // (already colored, already laid out) and emitted verbatim, so reflow neither + // wraps it nor overrides its colors with a role style. + roleBanner +) + +// transcriptBlock is one rendered turn in the transcript. text is the raw +// (unstyled, unwrapped) message body; the role selects the theme style and any +// prefix applied at render time. For roleTool blocks text is unused and card +// points at the live tool card (#389); the pointer lets a later toolEndMsg / +// Ctrl+O mutate the card in place and have it re-render on the next reflow. +type transcriptBlock struct { + role blockRole + text string + card *toolCard +} + +// transcript is the scrolling message log. It wraps a viewport.Model and keeps +// the source blocks so it can re-flow on width changes. activeAssistant indexes +// the assistant block currently receiving streaming deltas, or -1 when no turn +// is streaming. +type transcript struct { + vp viewport.Model + theme Theme + + // totalWidth is the full width the transcript may occupy (terminal columns + // minus any chrome the model reserves). width (below) is the content width + // the blocks actually wrap to: it equals totalWidth when the content fits, or + // totalWidth-1 when it overflows and a scrollbar column must be held back. + // reflow recomputes width from totalWidth on every content change, so the bar + // column appears/disappears correctly even as a run streams in new lines. + totalWidth int + + // width is the content width (terminal columns) the blocks wrap to. It is + // separate from the viewport's own width so reflow measurements stay stable + // even before the first size message. + width int + + blocks []transcriptBlock + activeAssistant int + + // follow is the stick-to-bottom intent: while true, every reflow snaps the + // viewport to the newest line so streamed output stays visible. It is set + // when the user submits a turn and cleared when they scroll up to read + // history (re-armed when they scroll back to the bottom). Tracking intent + // explicitly — rather than sampling viewport.AtBottom() inside reflow — keeps + // auto-scroll correct across height changes (setSize resizes the viewport + // before reflow runs, which would make an AtBottom() sample read false). + follow bool +} + +// newTranscript builds an empty transcript with the given theme. The viewport +// starts zero-sized; the model drives setSize from the first tea.WindowSizeMsg. +func newTranscript(theme Theme) transcript { + vp := viewport.New() + return transcript{ + vp: vp, + theme: theme, + activeAssistant: -1, + } +} + +// setSize resizes the transcript's viewport and re-flows the blocks to the new +// width. A non-positive dimension is clamped to zero so the viewport never sees +// a negative extent. width is the total space available; reflow decides whether +// to spend one column on the scrollbar based on whether the content overflows. +func (t *transcript) setSize(width, height int) { + if width < 0 { + width = 0 + } + if height < 0 { + height = 0 + } + t.totalWidth = width + t.vp.SetHeight(height) + t.reflow() +} + +// addUser appends a user turn and closes any streaming assistant block, then +// re-flows. Submitting a prompt is an explicit action where the user always +// wants to see their new turn and the response that follows, so it re-arms +// follow: the viewport snaps to the bottom even if the user had scrolled up +// (e.g. reading the startup banner) — otherwise the streamed reply would +// accumulate off-screen and look like nothing happened. Subsequent streaming +// deltas keep the bottom via follow, which the user can pause by scrolling up. +func (t *transcript) addUser(text string) { + t.blocks = append(t.blocks, transcriptBlock{role: roleUser, text: text}) + t.activeAssistant = -1 + t.follow = true + t.reflow() +} + +// addSystem appends a system / meta notice (used for run lifecycle and other +// inline notes). +func (t *transcript) addSystem(text string) { + t.blocks = append(t.blocks, transcriptBlock{role: roleSystem, text: text}) + t.reflow() +} + +// addBanner appends a pre-rendered splash block (startup logo + config). It is +// emitted verbatim by renderBlock, so its colors and horizontal layout survive +// reflow untouched. +func (t *transcript) addBanner(text string) { + t.blocks = append(t.blocks, transcriptBlock{role: roleBanner, text: text}) + t.reflow() +} + +// addToolCard appends a rich tool-call card (#389) as an ordered block so it +// renders inline in the transcript. The card is held by pointer, so a later +// state change (toolEndMsg) or expand toggle (Ctrl+O) followed by reflow +// re-renders it in place. +func (t *transcript) addToolCard(c *toolCard) { + t.blocks = append(t.blocks, transcriptBlock{role: roleTool, card: c}) + t.reflow() +} + +// appendDelta grows the current assistant block by delta, creating the block on +// the first delta of a turn. The re-flow auto-sticks to the bottom when the user +// has not scrolled up. +func (t *transcript) appendDelta(delta string) { + if t.activeAssistant < 0 { + t.blocks = append(t.blocks, transcriptBlock{role: roleAssistant}) + t.activeAssistant = len(t.blocks) - 1 + } + t.blocks[t.activeAssistant].text += delta + t.reflow() +} + +// finalizeTurn closes the streaming assistant block. When the final message +// carries text it becomes the block's authoritative body (covering turns that +// arrive without incremental deltas); otherwise the accumulated deltas stand. +func (t *transcript) finalizeTurn(msg agentcore.AssistantMessage) { + text := agentcore.ContentToText(msg.Content) + if t.activeAssistant >= 0 { + if text != "" { + t.blocks[t.activeAssistant].text = text + } + } else if text != "" { + t.blocks = append(t.blocks, transcriptBlock{role: roleAssistant, text: text}) + } + t.activeAssistant = -1 + t.reflow() +} + +// update forwards a message (typically a key press or scroll) to the viewport so +// PgUp/PgDn/arrow scrolling works, then re-syncs the follow intent: scrolling up +// off the bottom pauses auto-scroll, and scrolling back to the bottom re-arms it. +func (t *transcript) update(msg tea.Msg) tea.Cmd { + var cmd tea.Cmd + t.vp, cmd = t.vp.Update(msg) + t.follow = t.vp.AtBottom() + return cmd +} + +// scrollToRow positions the viewport so the scrollbar thumb aligns with the +// given viewport row y (0-based). It is the inverse of the thumb-position math +// in scrollbar(): pressing or dragging on row y maps that row to the matching +// scroll offset, so clicking the gutter jumps there and dragging the thumb +// tracks the cursor. It is a no-op when the content fits (nothing to scroll). +func (t *transcript) scrollToRow(y int) { + h := t.vp.Height() + if h <= 0 { + return + } + total := t.vp.TotalLineCount() + if total <= h { + return + } + thumb := h * h / total + if thumb < 2 { + thumb = 2 + } + if thumb > h { + thumb = h + } + span := h - thumb // rows the thumb top can occupy + if span <= 0 { + return + } + // Center the grab on the thumb: aim its top at y minus half its body so the + // cursor sits roughly mid-thumb, then clamp into the track. + top := y - thumb/2 + if top < 0 { + top = 0 + } + if top > span { + top = span + } + maxOff := total - h + t.vp.SetYOffset(top * maxOff / span) + t.follow = t.vp.AtBottom() +} + +// viewportHeight reports the number of visible transcript rows, so the model can +// tell whether a mouse Y falls within the scrollable region. +func (t transcript) viewportHeight() int { return t.vp.Height() } + +// overflowing reports whether the transcript has more content than fits in the +// viewport, i.e. there is history to scroll. relayout uses this to reserve the +// scrollbar column only when scrolling is possible, and view uses it to decide +// whether to attach the thumb at all. +func (t transcript) overflowing() bool { + return t.vp.Height() > 0 && t.vp.TotalLineCount() > t.vp.Height() +} + +// view renders the current visible slice of the transcript. When the content +// overflows the viewport a one-column vertical scrollbar is drawn down the right +// edge (FR-10): each viewport row is normalized to exactly the content width +// before the scrollbar cell is appended, so the bar sits flush against the +// terminal's right edge and a dangling SGR from Markdown rendering can never +// bleed into (and hide) the bar column. When everything fits there is nothing to +// scroll, so no bar is drawn and the viewport uses the full width (relayout +// releases the reserved column in that case). +func (t transcript) view() string { + if !t.overflowing() { + return t.vp.View() + } + + bar := strings.Split(t.scrollbar(), "\n") + body := strings.Split(t.vp.View(), "\n") + + // Fit every body line to exactly t.width columns (ANSI-aware pad/truncate), + // terminating any open style so the bar cell renders on a clean slate. + fit := lipgloss.NewStyle().Width(t.width).MaxWidth(t.width) + + var b strings.Builder + for i := 0; i < len(bar); i++ { + if i > 0 { + b.WriteByte('\n') + } + line := "" + if i < len(body) { + line = body[i] + } + if t.width > 0 { + b.WriteString(fit.Render(line)) + } + b.WriteString(bar[i]) + } + return b.String() +} + +// scrollbar renders the one-column vertical scrollbar the height of the +// viewport. A proportional thumb marks the visible window and its position marks +// the scroll offset, so scrolling up through history moves the thumb; the +// remaining rows draw a thin groove (│). The thumb is drawn as a capsule like +// the macOS system scrollbar: a lower-half block ▄ caps the top and an upper-half +// block ▀ caps the bottom (their filled halves sit on the inner edges so the +// outer ends taper to rounded), with the full block █ filling the body rows +// between the caps. The thumb is never shorter than three rows, so the capsule +// always shows a body between its two rounded caps rather than collapsing to a +// flat blob. When the content fits (no overflow) the capsule fills the full +// height. +func (t transcript) scrollbar() string { + h := t.vp.Height() + if h <= 0 { + return "" + } + total := t.vp.TotalLineCount() + thumb := h + pos := 0 + if total > h { + thumb = h * h / total + // Keep the capsule shape (rounded cap + body + rounded cap) by never + // letting the thumb shrink below three rows; clamp down to the viewport + // height when it is shorter than that. + if thumb < 3 { + thumb = 3 + } + if thumb > h { + thumb = h + } + maxOff := total - h + off := t.vp.YOffset() + if off > maxOff { + off = maxOff + } + if maxOff > 0 { + pos = off * (h - thumb) / maxOff + } + } + var b strings.Builder + for i := 0; i < h; i++ { + if i > 0 { + b.WriteByte('\n') + } + switch { + case i < pos || i >= pos+thumb: + b.WriteString(t.theme.ScrollTrack.Render("│")) + case thumb >= 2 && i == pos: + b.WriteString(t.theme.ScrollThumb.Render("▄")) + case thumb >= 2 && i == pos+thumb-1: + b.WriteString(t.theme.ScrollThumb.Render("▀")) + default: + b.WriteString(t.theme.ScrollThumb.Render("█")) + } + } + return b.String() +} + +// reflow re-renders every block to the current width and pushes the joined +// content into the viewport. When the follow intent is set it snaps to the +// bottom so new content auto-scrolls; otherwise the offset is preserved so +// reading history is not interrupted. follow is tracked in update/scrollToRow +// (user scroll) and addUser (new turn) rather than sampled here, because setSize +// resizes the viewport before reflow runs and an AtBottom() sample would misread. +// +// Width is decided here rather than in setSize so it stays correct as a run +// streams in new lines (which reach reflow via appendDelta/finalizeTurn, not +// setSize): the blocks are first laid out at the full width, and only if that +// overflows the viewport is one column handed back to the scrollbar and the +// blocks re-laid at totalWidth-1. When the content fits, the transcript keeps +// the full width and view() draws no bar. +func (t *transcript) reflow() { + t.width = t.totalWidth + t.vp.SetWidth(t.width) + t.vp.SetContent(t.renderAll()) + + // A narrower width never reduces the line count, so if the full-width layout + // already overflows it still overflows at totalWidth-1: reserve the scrollbar + // column and re-lay the blocks so the body never sits under the bar. + if t.totalWidth > 0 && t.vp.TotalLineCount() > t.vp.Height() { + t.width = t.totalWidth - 1 + t.vp.SetWidth(t.width) + t.vp.SetContent(t.renderAll()) + } + + if t.follow { + t.vp.GotoBottom() + } +} + +// renderAll joins every block, rendered to the current content width, into the +// transcript body string. Consecutive turns are separated by a blank line before +// a new user turn so requests read as visually distinct. +func (t *transcript) renderAll() string { + var b strings.Builder + for i, blk := range t.blocks { + if i > 0 { + b.WriteByte('\n') + if blk.role == roleUser { + b.WriteByte('\n') + } + } + b.WriteString(t.renderBlock(blk, i == t.activeAssistant)) + } + return b.String() +} + +// renderBlock wraps a block's text to the content width and applies the role's +// theme style. Wrapping happens on the raw text (measured in display columns via +// WrapToWidth) before styling so ANSI escapes never confuse the width math and +// no double-width rune is split. A finalized assistant block is rendered as +// Markdown (fix #3, mirroring the REPL's turn-end render); the still-streaming +// block (streaming==true) stays plain text because Markdown can only be laid out +// once the whole block is known. +func (t transcript) renderBlock(blk transcriptBlock, streaming bool) string { + if blk.role == roleTool && blk.card != nil { + return blk.card.render(t.theme, t.width) + } + switch blk.role { + case roleBanner: + return blk.text + case roleUser: + return t.theme.User.Render(WrapToWidth(blk.text, t.width)) + case roleSystem: + return t.theme.System.Render(WrapToWidth(blk.text, t.width)) + default: + if streaming { + return t.theme.Assistant.Render(WrapToWidth(blk.text, t.width)) + } + return renderMarkdown(blk.text, t.width) + } +} diff --git a/pigo/internal/cli/tui/transcript_test.go b/pigo/internal/cli/tui/transcript_test.go new file mode 100644 index 0000000..3035109 --- /dev/null +++ b/pigo/internal/cli/tui/transcript_test.go @@ -0,0 +1,324 @@ +package tui + +import ( + "regexp" + "strings" + "testing" + + tea "charm.land/bubbletea/v2" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/cli/ui" +) + +// ansiRE strips SGR escape sequences so tests can inspect the raw text the +// transcript stored, independent of the theme's coloring. +var ansiRE = regexp.MustCompile("\x1b\\[[0-9;]*m") + +func stripANSI(s string) string { return ansiRE.ReplaceAllString(s, "") } + +// apply runs one Update tick and returns the concrete Model, failing on an +// unexpected model type. It keeps the streaming tests terse. +func apply(t *testing.T, m tea.Model, msg tea.Msg) Model { + t.Helper() + next, _ := m.Update(msg) + got, ok := next.(Model) + if !ok { + t.Fatalf("Update returned %T, want tui.Model", next) + } + return got +} + +// TestTranscriptStreamingConcat feeds a run of text deltas then a turn end and +// asserts the assistant block accumulates the deltas in order and the joined +// text is rendered in the View. +func TestTranscriptStreamingConcat(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + m = apply(t, m, textDeltaMsg{delta: "Hello "}) + m = apply(t, m, textDeltaMsg{delta: "world"}) + m = apply(t, m, turnEndMsg{msg: agentcore.AssistantMessage{ + Content: agentcore.ContentList{agentcore.NewTextContent("Hello world")}, + }}) + + if n := len(m.transcript.blocks); n != 1 { + t.Fatalf("block count = %d, want 1 assistant block", n) + } + if got := m.transcript.blocks[0]; got.role != roleAssistant || got.text != "Hello world" { + t.Errorf("assistant block = %+v, want role assistant text %q", got, "Hello world") + } + if content := stripANSI(m.View().Content); !strings.Contains(content, "Hello world") { + t.Errorf("rendered View missing streamed text; got:\n%s", content) + } + // The turn was finalized, so a fresh delta starts a NEW assistant block. + if m.transcript.activeAssistant != -1 { + t.Errorf("activeAssistant = %d after turn end, want -1", m.transcript.activeAssistant) + } +} + +// TestTranscriptSurfacesTurnError verifies a turn that ends with stopReason +// error surfaces the provider's error message as a system block rather than +// finalizing an empty turn and returning silently to the prompt. The loop +// delivers request failures (e.g. a 4xx) this way — as a terminal assistant +// message via TurnEndEvent, not as the run's result error — so without the +// StopReason check in the turnEndMsg handler the TUI would show nothing at all. +func TestTranscriptSurfacesTurnError(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + m = apply(t, m, turnEndMsg{msg: agentcore.AssistantMessage{ + StopReason: agentcore.StopReasonError, + ErrorMessage: "upstream 401: 无效的令牌", + }}) + + var sys string + for _, b := range m.transcript.blocks { + if b.role == roleSystem { + sys = b.text + } + } + if !strings.Contains(sys, "error:") || !strings.Contains(sys, "upstream 401: 无效的令牌") { + t.Errorf("turn error not surfaced; system block = %q", sys) + } + if content := stripANSI(m.View().Content); !strings.Contains(content, "upstream 401") { + t.Errorf("rendered View missing the surfaced error; got:\n%s", content) + } +} + +// TestTranscriptSurfacesAbortedTurn verifies a turn that ends with stopReason +// aborted is flagged rather than returning silently. +func TestTranscriptSurfacesAbortedTurn(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + m = apply(t, m, turnEndMsg{msg: agentcore.AssistantMessage{ + StopReason: agentcore.StopReasonAborted, + }}) + + var sys string + for _, b := range m.transcript.blocks { + if b.role == roleSystem { + sys = b.text + } + } + if !strings.Contains(sys, "aborted") { + t.Errorf("aborted turn not surfaced; system block = %q", sys) + } +} + +// TestTranscriptNotesEmptyResponse verifies a clean end_turn that produced no +// content and no tool results is flagged with a note (with a provider-mismatch +// hint) instead of returning silently to the prompt — the shape produced when an +// endpoint accepts the request with a 200 but returns nothing decodable. +func TestTranscriptNotesEmptyResponse(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + m = apply(t, m, turnEndMsg{msg: agentcore.AssistantMessage{ + StopReason: agentcore.StopReasonEndTurn, + }}) + + var sys string + for _, b := range m.transcript.blocks { + if b.role == roleSystem { + sys = b.text + } + } + if !strings.Contains(sys, "empty response from the model") { + t.Errorf("empty response not flagged; system block = %q", sys) + } +} + +// TestTranscriptCleanTurnNoNote verifies a normal turn with content does NOT add +// a spurious error/empty system note. +func TestTranscriptCleanTurnNoNote(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 40, Height: 12}) + + m = apply(t, m, turnEndMsg{msg: agentcore.AssistantMessage{ + StopReason: agentcore.StopReasonEndTurn, + Content: agentcore.ContentList{agentcore.NewTextContent("the answer")}, + }}) + + for _, b := range m.transcript.blocks { + if b.role == roleSystem { + t.Errorf("clean turn should add no system note, got %q", b.text) + } + } +} + +// TestTranscriptAutoStick verifies the stick-to-bottom rule: while the viewport +// is at the bottom, new content keeps it pinned there; once the user scrolls up, +// streamed content no longer forces a jump to the bottom — but submitting a new +// turn (addUser) re-arms follow and snaps back to the newest output. +func TestTranscriptAutoStick(t *testing.T) { + tr := newTranscript(DefaultTheme()) + tr.setSize(20, 3) // 3 visible rows + + for i := 0; i < 6; i++ { + tr.addUser("line") + } + if !tr.vp.AtBottom() { + t.Fatal("transcript should stick to the bottom while at the bottom") + } + + // More content while pinned keeps it pinned. + tr.addUser("more") + if !tr.vp.AtBottom() { + t.Fatal("new content should keep a bottom-pinned transcript at the bottom") + } + + // Simulate a user scroll-up through the viewport's key handling. + tr.update(tea.KeyPressMsg{Code: tea.KeyUp}) + if tr.vp.AtBottom() { + t.Fatal("scrolling up should move the viewport off the bottom") + } + + // Streamed content arriving while scrolled up must NOT yank the view back to + // the bottom — the user is reading history. + tr.appendDelta("streamed while reading history\nsecond line\nthird line") + if tr.vp.AtBottom() { + t.Error("auto-stick should stay paused after the user scrolls up") + } + + // Submitting a new turn is an explicit action: it re-arms follow and snaps + // back to the newest output so the reply is never left off-screen. + tr.addUser("a brand new prompt") + if !tr.vp.AtBottom() { + t.Error("submitting a new turn should re-arm auto-scroll to the bottom") + } +} + +// TestTranscriptCJKWrap feeds a long CJK line into a narrow transcript and +// asserts every wrapped line fits the width in display columns (not bytes) and +// that no rune was dropped or split. +func TestTranscriptCJKWrap(t *testing.T) { + const width = 10 + tr := newTranscript(DefaultTheme()) + tr.setSize(width, 20) + + line := strings.Repeat("你好世界", 5) // 20 CJK runes = 40 display columns + tr.addUser(line) + + content := tr.vp.GetContent() + lines := strings.Split(content, "\n") + if len(lines) < 2 { + t.Fatalf("expected the CJK line to wrap onto multiple rows, got %d line(s)", len(lines)) + } + for i, ln := range lines { + if w := ui.Width(ln); w > width { + t.Errorf("wrapped line %d width = %d columns, want <= %d: %q", i, w, width, stripANSI(ln)) + } + } + // No rune was cut or dropped: every source rune survives the wrap. + if got := strings.Count(stripANSI(content), "你"); got != 5 { + t.Errorf("counted %d 你 runes after wrap, want 5", got) + } +} + +// TestTranscriptScrollbar verifies the scrollbar policy: the gutter is hidden +// while the content fits (nothing to scroll) and appears only once the content +// overflows. When overflowing, a rounded pill thumb (body "█" with half-block +// caps "▄"/"▀") sits alongside the thin groove "│". +func TestTranscriptScrollbar(t *testing.T) { + tr := newTranscript(DefaultTheme()) + tr.setSize(20, 4) // 4 visible rows + + // Two short lines fit in 4 rows: no scrollbar at all — no thumb, no groove. + tr.addUser("one") + tr.addUser("two") + if tr.overflowing() { + t.Fatal("transcript should not overflow while content fits") + } + fit := stripANSI(tr.view()) + if strings.ContainsAny(fit, "█▄▀│") { + t.Errorf("expected no scrollbar glyphs while content fits; got:\n%q", fit) + } + + // Enough lines to exceed 4 rows: now it overflows, thumb shrinks and the + // groove appears. + for i := 0; i < 10; i++ { + tr.addUser("line") + } + if !tr.overflowing() { + t.Fatal("transcript should overflow once content exceeds the viewport") + } + view := stripANSI(tr.view()) + if !strings.Contains(view, "▄") || !strings.Contains(view, "▀") { + t.Errorf("expected a rounded pill thumb (▄ top, ▀ bottom) while overflowing; got:\n%q", view) + } + if !strings.Contains(view, "│") { + t.Errorf("expected a groove │ while overflowing; got:\n%q", view) + } + if strings.Contains(view, "░") { + t.Errorf("scrollbar no longer uses the shaded track ░; got:\n%q", view) + } +} + +// TestTranscriptScrollToRow checks the click/drag mapping: pressing the top of +// the gutter scrolls to the top, the bottom scrolls to the bottom, and it is a +// no-op when the content fits. +func TestTranscriptScrollToRow(t *testing.T) { + tr := newTranscript(DefaultTheme()) + tr.setSize(20, 4) + + // Content fits: dragging must not move a non-scrollable viewport. + tr.addUser("only line") + tr.scrollToRow(3) + if tr.vp.YOffset() != 0 { + t.Errorf("scrollToRow on non-overflowing viewport moved offset to %d, want 0", tr.vp.YOffset()) + } + + for i := 0; i < 20; i++ { + tr.addUser("line") + } + if !tr.overflowing() { + t.Fatal("expected overflow after filling the transcript") + } + + tr.scrollToRow(0) + if !tr.vp.AtTop() { + t.Errorf("dragging to row 0 should scroll to the top; YOffset=%d", tr.vp.YOffset()) + } + + tr.scrollToRow(tr.viewportHeight() - 1) + if !tr.vp.AtBottom() { + t.Errorf("dragging to the last row should scroll to the bottom; YOffset=%d", tr.vp.YOffset()) + } +} + +// TestModelScrollbarDrag drives the model with mouse press/motion/release on the +// scrollbar column and asserts the drag state toggles and the viewport scrolls. +func TestModelScrollbarDrag(t *testing.T) { + m := apply(t, NewModel(Options{}), tea.WindowSizeMsg{Width: 30, Height: 8}) + for i := 0; i < 40; i++ { + m.transcript.addUser("line") + } + if !m.transcript.overflowing() { + t.Fatal("expected the transcript to overflow") + } + + col := m.width - 1 + // Press at the top of the gutter: drag begins and the view jumps to the top. + m = apply(t, m, tea.MouseClickMsg{X: col, Y: 0, Button: tea.MouseLeft}) + if !m.draggingScrollbar { + t.Fatal("left press on the scrollbar column should start dragging") + } + if !m.transcript.vp.AtTop() { + t.Errorf("press at row 0 should scroll to top; YOffset=%d", m.transcript.vp.YOffset()) + } + + // Motion to the bottom row while held drags the thumb down. + m = apply(t, m, tea.MouseMotionMsg{X: col, Y: m.transcript.viewportHeight() - 1, Button: tea.MouseLeft}) + if !m.transcript.vp.AtBottom() { + t.Errorf("motion to the last row while dragging should scroll to bottom; YOffset=%d", m.transcript.vp.YOffset()) + } + + // Release ends the drag. + m = apply(t, m, tea.MouseReleaseMsg{X: col, Y: 3, Button: tea.MouseLeft}) + if m.draggingScrollbar { + t.Error("release should end the scrollbar drag") + } + + // A press away from the gutter column must not start a drag. + m = apply(t, m, tea.MouseClickMsg{X: 0, Y: 0, Button: tea.MouseLeft}) + if m.draggingScrollbar { + t.Error("press off the scrollbar column should not start dragging") + } +} diff --git a/pigo/internal/cli/ui/color.go b/pigo/internal/cli/ui/color.go new file mode 100644 index 0000000..c9e2b77 --- /dev/null +++ b/pigo/internal/cli/ui/color.go @@ -0,0 +1,52 @@ +// Package ui holds the leaf terminal-UI helpers shared across the cmd/pigo and +// internal/cli subpackages: ANSI color gating (color.go), turn-end Markdown +// rendering (markdown.go), and prompt image-reference parsing (imageref.go). +// These were moved verbatim from cmd/pigo (US-002, #358) and exported so the +// repl, btw, status and goal layers style output through one owner. +package ui + +import "os" + +// ANSI SGR escape sequences used by the REPL. This is a handful of codes, not a +// general-purpose styling library. +const ( + Reset = "\033[0m" + Bold = "\033[1m" + Dim = "\033[2m" + Cyan = "\033[36m" + Green = "\033[32m" + Red = "\033[31m" + Yellow = "\033[33m" +) + +// StdoutIsTerminal reports whether stdout is an interactive terminal (not a +// pipe/file). It gates color output and is also used to decide print vs +// interactive mode. +func StdoutIsTerminal() bool { + fi, err := os.Stdout.Stat() + if err != nil { + return false + } + return fi.Mode()&os.ModeCharDevice != 0 +} + +// Enabled reports whether ANSI color should be emitted. Color is on only when +// stdout is an interactive terminal and NO_COLOR is unset (mirrors the +// https://no-color.org convention). This keeps piped/redirected output and CI +// logs free of escape codes. +func Enabled() bool { + if _, ok := os.LookupEnv("NO_COLOR"); ok { + return false + } + return StdoutIsTerminal() +} + +// Colorize wraps s in the given SGR code(s) and a reset when color is enabled, +// and returns s unchanged otherwise. Callers decide the code; an empty code +// returns s as-is so it is safe to call unconditionally. +func Colorize(enabled bool, code, s string) string { + if !enabled || code == "" { + return s + } + return code + s + Reset +} diff --git a/pigo/internal/cli/ui/color_test.go b/pigo/internal/cli/ui/color_test.go new file mode 100644 index 0000000..6f42dff --- /dev/null +++ b/pigo/internal/cli/ui/color_test.go @@ -0,0 +1,27 @@ +package ui + +import "testing" + +// TestColorizeGating verifies Colorize wraps text in SGR codes only when +// enabled, returns text unchanged when disabled, and treats an empty code as a +// no-op regardless of the enabled flag. +func TestColorizeGating(t *testing.T) { + if got := Colorize(true, Cyan, "/help"); got != Cyan+"/help"+Reset { + t.Errorf("enabled: got %q", got) + } + if got := Colorize(false, Cyan, "/help"); got != "/help" { + t.Errorf("disabled should be plain, got %q", got) + } + if got := Colorize(true, "", "/help"); got != "/help" { + t.Errorf("empty code should be plain, got %q", got) + } +} + +// TestColorEnabledRespectsNoColor verifies NO_COLOR forces color off even on a +// terminal (mirrors https://no-color.org). +func TestColorEnabledRespectsNoColor(t *testing.T) { + t.Setenv("NO_COLOR", "1") + if Enabled() { + t.Error("NO_COLOR set: Enabled must be false") + } +} diff --git a/pigo/internal/cli/ui/imageref.go b/pigo/internal/cli/ui/imageref.go new file mode 100644 index 0000000..af3ddc1 --- /dev/null +++ b/pigo/internal/cli/ui/imageref.go @@ -0,0 +1,119 @@ +// This file implements the local-image input syntax for prompts (US-010/#126). +// A prompt may reference local image files with either `@image:` or the +// Markdown image form `![alt]()`. Each reference is read from disk, +// base64-encoded, and attached to the user message as an agentcore.ImageContent +// block so a multimodal model can see it. The remaining (non-reference) text is +// kept as a TextContent block. References that cannot be read are reported as an +// error rather than silently dropped, so the user knows the image was not sent. +package ui + +import ( + "encoding/base64" + "fmt" + "net/http" + "os" + "regexp" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// imageRefPattern matches the two supported image-reference syntaxes: +// - @image: (path runs to the next whitespace) +// - ![alt]() (Markdown image; alt text is ignored) +// +// The path in the Markdown form may contain spaces; the @image form may not +// (whitespace terminates it), matching the convention that @image is a bare +// token while the Markdown form is explicitly delimited by parentheses. +var imageRefPattern = regexp.MustCompile(`!\[[^\]]*\]\(([^)]+)\)|@image:(\S+)`) + +// BuildUserContent turns a raw prompt line into a content list. When the prompt +// contains no image references it returns a single TextContent (the common +// case). Otherwise it returns the interleaved text and image blocks, reading and +// base64-encoding each referenced file. It returns an error if any referenced +// file cannot be read, so the caller can surface it instead of sending a prompt +// that silently omits the image. +func BuildUserContent(prompt string) (agentcore.ContentList, error) { + locs := imageRefPattern.FindAllStringSubmatchIndex(prompt, -1) + if len(locs) == 0 { + return agentcore.ContentList{agentcore.NewTextContent(prompt)}, nil + } + + var content agentcore.ContentList + // addText appends a text block for the given [lo,hi) slice of prompt, + // trimming surrounding whitespace and skipping empties so we don't emit + // blank text blocks around the references. + addText := func(lo, hi int) { + if lo >= hi { + return + } + if t := strings.TrimSpace(prompt[lo:hi]); t != "" { + content = append(content, agentcore.NewTextContent(t)) + } + } + + prev := 0 + for _, loc := range locs { + start, end := loc[0], loc[1] + addText(prev, start) + // Group 1 is the Markdown path; group 2 is the @image path. Exactly one + // is set per match. + var path string + if loc[2] >= 0 { + path = prompt[loc[2]:loc[3]] + } else if loc[4] >= 0 { + path = prompt[loc[4]:loc[5]] + } + img, err := loadImageContent(strings.TrimSpace(path)) + if err != nil { + return nil, err + } + content = append(content, img) + prev = end + } + addText(prev, len(prompt)) + + if len(content) == 0 { + content = agentcore.ContentList{agentcore.NewTextContent("")} + } + return content, nil +} + +// loadImageContent reads an image file and returns an ImageContent with the +// data base64-encoded and the mime type sniffed from the file extension (with a +// content-sniff fallback). It errors if the file cannot be read or is not a +// recognizable image type. +func loadImageContent(path string) (agentcore.ImageContent, error) { + data, err := os.ReadFile(path) + if err != nil { + return agentcore.ImageContent{}, fmt.Errorf("read image %q: %w", path, err) + } + mime := mimeFromPath(path) + if mime == "" { + // Fall back to content sniffing when the extension is unknown. + mime = http.DetectContentType(data) + } + if !strings.HasPrefix(mime, "image/") { + return agentcore.ImageContent{}, fmt.Errorf("%q is not a recognized image (detected %q)", path, mime) + } + enc := base64.StdEncoding.EncodeToString(data) + return agentcore.NewImageContent(enc, mime), nil +} + +// mimeFromPath maps a file extension to an image mime type. It returns "" for +// unknown extensions so the caller can fall back to content sniffing. +func mimeFromPath(path string) string { + lower := strings.ToLower(path) + switch { + case strings.HasSuffix(lower, ".png"): + return "image/png" + case strings.HasSuffix(lower, ".jpg"), strings.HasSuffix(lower, ".jpeg"): + return "image/jpeg" + case strings.HasSuffix(lower, ".gif"): + return "image/gif" + case strings.HasSuffix(lower, ".webp"): + return "image/webp" + default: + return "" + } +} diff --git a/pigo/internal/cli/ui/imageref_test.go b/pigo/internal/cli/ui/imageref_test.go new file mode 100644 index 0000000..9046a5a --- /dev/null +++ b/pigo/internal/cli/ui/imageref_test.go @@ -0,0 +1,122 @@ +// Tests for the prompt image-reference syntax (US-010, #126): @image: and +// the Markdown ![alt]() form. They write a tiny real PNG to a temp file so +// loadImageContent exercises the real read + base64 + mime path, and assert that +// a missing file is an error (not a silently dropped image). +package ui + +import ( + "encoding/base64" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// pngBytes is a minimal 1x1 PNG (valid signature + IHDR) so mime detection and +// the image/ prefix check pass without pulling in an image library. +var pngBytes = []byte{ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, + 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, + 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, + 0x08, 0x06, 0x00, 0x00, 0x00, 0x1F, 0x15, 0xC4, + 0x89, +} + +func writeTempPNG(t *testing.T) string { + t.Helper() + dir := t.TempDir() + path := filepath.Join(dir, "pixel.png") + if err := os.WriteFile(path, pngBytes, 0o644); err != nil { + t.Fatalf("write temp png: %v", err) + } + return path +} + +// TestBuildUserContentNoImage: a plain prompt yields a single text block. +func TestBuildUserContentNoImage(t *testing.T) { + content, err := BuildUserContent("just some text") + if err != nil { + t.Fatalf("BuildUserContent: %v", err) + } + if len(content) != 1 { + t.Fatalf("content len = %d, want 1", len(content)) + } + if tc, ok := content[0].(agentcore.TextContent); !ok || tc.Text != "just some text" { + t.Errorf("content[0] = %#v, want text \"just some text\"", content[0]) + } +} + +// TestBuildUserContentAtImageSyntax: @image: attaches an ImageContent and +// keeps the surrounding text. +func TestBuildUserContentAtImageSyntax(t *testing.T) { + path := writeTempPNG(t) + content, err := BuildUserContent("look at @image:" + path + " please") + if err != nil { + t.Fatalf("BuildUserContent: %v", err) + } + assertTextThenImage(t, content, "look at", "image/png") +} + +// TestBuildUserContentMarkdownSyntax: ![alt](path) attaches an ImageContent. +func TestBuildUserContentMarkdownSyntax(t *testing.T) { + path := writeTempPNG(t) + content, err := BuildUserContent("before ![a pixel](" + path + ") after") + if err != nil { + t.Fatalf("BuildUserContent: %v", err) + } + // text "before", image, text "after" + if len(content) != 3 { + t.Fatalf("content len = %d, want 3", len(content)) + } + if _, ok := content[1].(agentcore.ImageContent); !ok { + t.Errorf("content[1] = %T, want ImageContent", content[1]) + } +} + +// TestBuildUserContentMissingFile: a missing referenced file is an error. +func TestBuildUserContentMissingFile(t *testing.T) { + _, err := BuildUserContent("@image:/no/such/file.png") + if err == nil { + t.Fatal("BuildUserContent accepted a missing image file, want error") + } +} + +// TestLoadImageContentBase64: the encoded data must match the file bytes. +func TestLoadImageContentBase64(t *testing.T) { + path := writeTempPNG(t) + img, err := loadImageContent(path) + if err != nil { + t.Fatalf("loadImageContent: %v", err) + } + if img.MimeType != "image/png" { + t.Errorf("mime = %q, want image/png", img.MimeType) + } + if want := base64.StdEncoding.EncodeToString(pngBytes); img.Data != want { + t.Errorf("data mismatch") + } +} + +func assertTextThenImage(t *testing.T, content agentcore.ContentList, wantText, wantMime string) { + t.Helper() + var sawText, sawImage bool + for _, c := range content { + switch b := c.(type) { + case agentcore.TextContent: + if strings.Contains(b.Text, wantText) { + sawText = true + } + case agentcore.ImageContent: + if b.MimeType == wantMime { + sawImage = true + } + } + } + if !sawText { + t.Errorf("no text block containing %q in %#v", wantText, content) + } + if !sawImage { + t.Errorf("no image block with mime %q in %#v", wantMime, content) + } +} diff --git a/pigo/internal/cli/ui/markdown.go b/pigo/internal/cli/ui/markdown.go new file mode 100644 index 0000000..0f8e8b1 --- /dev/null +++ b/pigo/internal/cli/ui/markdown.go @@ -0,0 +1,72 @@ +// This file renders assistant replies as Markdown for interactive terminals. +// The line-oriented REPL streams model text token-by-token, but Markdown can +// only be laid out once the whole block is known (a table or fenced code span +// needs its full extent). So rendering is a turn-end concern: the caller buffers +// the streamed text and calls RenderMarkdown once the assistant turn closes. +// +// Rendering is gated exactly like color (Enabled): only an interactive, +// NO_COLOR-unset stdout gets styled output. Pipes, files, CI, and tests receive +// the raw Markdown source unchanged, so machine consumers and golden tests are +// unaffected. Any renderer failure also falls back to the raw source — pretty +// output is never allowed to lose content. +package ui + +import ( + "strings" + "sync" + + "github.com/charmbracelet/glamour" +) + +// mdRenderer is the lazily-built glamour renderer. Building it parses a style +// and compiles a chroma lexer set, so it is created once and reused across +// turns. A build failure leaves it nil, degrading to raw output. +var ( + mdOnce sync.Once + mdRenderer *glamour.TermRenderer +) + +// initMarkdown builds the shared renderer on first use. It uses glamour's +// auto style, which follows the terminal's dark/light background. +// +// WithWordWrap(0) disables glamour's hard word-wrap. That matters: with a fixed +// wrap width glamour pads every line with trailing-space background cells out to +// the full column count, so a three-line reply balloons into kilobytes of ANSI +// noise (measured: ~8KB for a short block at width 100 vs. ~0.5KB unwrapped). +// Disabling the wrap lets the terminal soft-wrap long lines itself and keeps the +// rendered output tight — the REPL doesn't track terminal size anyway. +func initMarkdown() { + mdOnce.Do(func() { + r, err := glamour.NewTermRenderer( + glamour.WithAutoStyle(), + glamour.WithWordWrap(0), + ) + if err != nil { + return + } + mdRenderer = r + }) +} + +// RenderMarkdown returns src rendered as styled terminal Markdown when output +// is an interactive terminal, and src unchanged otherwise. A nil/broken +// renderer or a render error also returns src, so content is never dropped in +// favor of styling. The returned string carries its own trailing newline from +// glamour; callers should not add another. +func RenderMarkdown(src string) string { + if !Enabled() { + return src + } + if strings.TrimSpace(src) == "" { + return src + } + initMarkdown() + if mdRenderer == nil { + return src + } + out, err := mdRenderer.Render(src) + if err != nil { + return src + } + return out +} diff --git a/pigo/internal/cli/ui/markdown_test.go b/pigo/internal/cli/ui/markdown_test.go new file mode 100644 index 0000000..e34f94e --- /dev/null +++ b/pigo/internal/cli/ui/markdown_test.go @@ -0,0 +1,23 @@ +package ui + +import "testing" + +// In tests stdout is not a terminal, so Enabled() is false and RenderMarkdown +// must return the source verbatim — this is the contract that keeps piped +// output, CI logs, and golden tests free of ANSI escapes. +func TestRenderMarkdownRawWhenNotTerminal(t *testing.T) { + src := "# Heading\n\nSome **bold** text.\n" + if got := RenderMarkdown(src); got != src { + t.Fatalf("RenderMarkdown on non-terminal = %q, want raw source unchanged", got) + } +} + +// An empty (or whitespace-only) reply must pass through untouched so the caller +// never prints a stray rendered blank block. +func TestRenderMarkdownEmptyPassthrough(t *testing.T) { + for _, src := range []string{"", " ", "\n\t\n"} { + if got := RenderMarkdown(src); got != src { + t.Fatalf("RenderMarkdown(%q) = %q, want unchanged", src, got) + } + } +} diff --git a/pigo/internal/cli/ui/toolrender.go b/pigo/internal/cli/ui/toolrender.go new file mode 100644 index 0000000..553d349 --- /dev/null +++ b/pigo/internal/cli/ui/toolrender.go @@ -0,0 +1,58 @@ +// This file holds the compact tool-activity renderers shared by the REPL, the +// /goal autonomous loop, and /btw side threads: a tool call is shown as a green +// "→ tool:" line and a tool result as a green "← result:" (or red "← error:") +// line, with multi-line output collapsed to one line. The todo tool is the one +// exception — its result is printed in full so the live checklist stays visible. +package ui + +import ( + "fmt" + "io" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// RenderToolResult prints a tool result to out: the todo tool's result is shown +// in full (indented) so the live checklist stays visible; every other result is +// collapsed to a single "← result:"/"← error:" line. +func RenderToolResult(out io.Writer, tr agentcore.ToolResultMessage) { + text := agentcore.ContentToText(tr.Content) + color := Enabled() + if tr.ToolName == "todo" && !tr.IsError { + fmt.Fprintln(out, " "+Colorize(color, Green, "← todo:")) + for _, line := range strings.Split(text, "\n") { + fmt.Fprintf(out, " %s\n", line) + } + return + } + if tr.IsError { + fmt.Fprintf(out, " %s %s\n", Colorize(color, Red, "← error:"), OneLine(text)) + return + } + fmt.Fprintf(out, " %s %s\n", Colorize(color, Green, "← result:"), OneLine(text)) +} + +// ToolCallLabel renders a tool call as "name args" for the compact "→ tool:" +// status. Empty or "{}" arguments collapse to just the name. +func ToolCallLabel(c agentcore.ToolCallContent) string { + args := strings.TrimSpace(string(c.Arguments)) + if args == "" || args == "{}" { + return c.Name + } + return c.Name + " " + OneLine(args) +} + +// OneLine collapses a possibly multi-line string into a single trimmed line, +// truncating very long values, for the compact tool-activity statuses. +func OneLine(s string) string { + s = strings.TrimSpace(s) + if i := strings.IndexByte(s, '\n'); i >= 0 { + s = s[:i] + " …" + } + const max = 120 + if len(s) > max { + s = s[:max] + " …" + } + return s +} diff --git a/pigo/internal/cli/ui/width.go b/pigo/internal/cli/ui/width.go new file mode 100644 index 0000000..02e6e2a --- /dev/null +++ b/pigo/internal/cli/ui/width.go @@ -0,0 +1,16 @@ +package ui + +import "charm.land/lipgloss/v2" + +// Width reports the number of terminal cells the rendered string s occupies. It +// delegates to lipgloss v2's width measurement, which strips ANSI escape +// sequences and counts East Asian wide / fullwidth runes (CJK, emoji) as two +// columns. This is the single width primitive the TUI and REPL layers style +// through (per the tui-agent SPEC), so alignment stays consistent across the +// codebase instead of each caller hand-rolling its own East-Asian-width table. +// +// For multi-line input, Width returns the width of the widest line (lipgloss +// measures the bounding box), matching how the renderer lays text out. +func Width(s string) int { + return lipgloss.Width(s) +} diff --git a/pigo/internal/cli/ui/width_test.go b/pigo/internal/cli/ui/width_test.go new file mode 100644 index 0000000..9162706 --- /dev/null +++ b/pigo/internal/cli/ui/width_test.go @@ -0,0 +1,33 @@ +package ui + +import "testing" + +// TestWidthASCII pins that plain ASCII counts one cell per rune. +func TestWidthASCII(t *testing.T) { + if got := Width("hello"); got != 5 { + t.Errorf("Width(\"hello\") = %d, want 5", got) + } +} + +// TestWidthCJK pins that East Asian wide runes count as two cells, which is the +// whole reason we route width through lipgloss instead of len/RuneCount. +func TestWidthCJK(t *testing.T) { + // Three CJK ideographs = 6 cells. + if got := Width("达克克"); got != 6 { + t.Errorf("Width(CJK x3) = %d, want 6", got) + } + // Mixed: "a达" = 1 + 2 = 3 cells. + if got := Width("a达"); got != 3 { + t.Errorf("Width(\"a达\") = %d, want 3", got) + } +} + +// TestWidthStripsANSI pins that SGR escape sequences do not add to the width, so +// a colored string aligns the same as its plain form. +func TestWidthStripsANSI(t *testing.T) { + plain := "error" + colored := Cyan + plain + Reset + if got, want := Width(colored), Width(plain); got != want { + t.Errorf("Width(colored) = %d, want %d (ANSI must not count)", got, want) + } +} diff --git a/pigo/internal/clipboard/clipboard.go b/pigo/internal/clipboard/clipboard.go new file mode 100644 index 0000000..8b45db9 --- /dev/null +++ b/pigo/internal/clipboard/clipboard.go @@ -0,0 +1,74 @@ +// Package clipboard writes text to the system clipboard by shelling out to the +// platform's clipboard utility (US-009, #125). It deliberately avoids any cgo or +// third-party dependency: it probes for the standard command-line tools +// (pbcopy on macOS, wl-copy / xclip / xsel on Linux) and pipes text to the first +// one found. When no utility is available it reports ErrUnavailable so the +// caller can degrade gracefully (e.g. print the text instead). +package clipboard + +import ( + "errors" + "os/exec" + "runtime" + "strings" +) + +// ErrUnavailable is returned by Copy when no supported clipboard utility is +// found on the host, so the caller can fall back to printing the content. +var ErrUnavailable = errors.New("clipboard: no clipboard utility available") + +// candidate is a clipboard-writing command: the executable plus the args that +// make it read the payload from stdin. +type candidate struct { + name string + args []string +} + +// candidates returns the clipboard-write commands to try, in priority order, +// for the current OS. On macOS pbcopy is always present; on Linux the Wayland +// tool is preferred, then the two common X11 tools. +func candidates() []candidate { + switch runtime.GOOS { + case "darwin": + return []candidate{{name: "pbcopy"}} + case "windows": + return []candidate{{name: "clip"}} + default: // linux and other unixes + return []candidate{ + {name: "wl-copy"}, + {name: "xclip", args: []string{"-selection", "clipboard"}}, + {name: "xsel", args: []string{"--clipboard", "--input"}}, + } + } +} + +// Copy writes text to the system clipboard using the first available platform +// utility. It returns ErrUnavailable if none is found (so callers can fall back +// to printing), or the underlying exec error if a utility was found but failed. +func Copy(text string) error { + for _, c := range candidates() { + path, err := exec.LookPath(c.name) + if err != nil { + continue + } + cmd := exec.Command(path, c.args...) + cmd.Stdin = strings.NewReader(text) + if err := cmd.Run(); err != nil { + return err + } + return nil + } + return ErrUnavailable +} + +// Available reports whether a clipboard utility is present, without writing +// anything. Useful for a caller that wants to phrase its output differently when +// it knows the copy will fail. +func Available() bool { + for _, c := range candidates() { + if _, err := exec.LookPath(c.name); err == nil { + return true + } + } + return false +} diff --git a/pigo/internal/clipboard/clipboard_test.go b/pigo/internal/clipboard/clipboard_test.go new file mode 100644 index 0000000..e0010a3 --- /dev/null +++ b/pigo/internal/clipboard/clipboard_test.go @@ -0,0 +1,61 @@ +package clipboard + +// Tests for the clipboard helper (US-009, #125). Copy shells out to a platform +// utility, so we cannot assert the OS clipboard actually changed in a hermetic +// test. Instead we verify the graceful-degradation contract: with no utility on +// PATH, Copy returns ErrUnavailable (so the REPL can fall back to printing) and +// Available reports false. We control the environment by pointing PATH at an +// empty temp dir. + +import ( + "errors" + "os" + "testing" +) + +// TestCopyUnavailableWhenNoUtility verifies Copy returns ErrUnavailable and +// Available returns false when PATH holds no clipboard utility. This is the +// contract the REPL relies on to degrade to printing. +func TestCopyUnavailableWhenNoUtility(t *testing.T) { + // Point PATH at an empty dir so exec.LookPath finds no pbcopy/xclip/etc. + empty := t.TempDir() + t.Setenv("PATH", empty) + + if Available() { + t.Error("Available() = true with empty PATH, want false") + } + err := Copy("hello") + if !errors.Is(err, ErrUnavailable) { + t.Errorf("Copy err = %v, want ErrUnavailable", err) + } +} + +// TestCopyUsesUtilityOnPath verifies Copy invokes a discovered utility and +// succeeds when the utility exits 0. We plant a fake executable named after the +// current platform's first candidate on PATH and confirm Copy returns nil. +func TestCopyUsesUtilityOnPath(t *testing.T) { + cands := candidates() + if len(cands) == 0 { + t.Skip("no clipboard candidates for this platform") + } + dir := t.TempDir() + name := cands[0].name + if os.PathSeparator == '\\' { + t.Skip("fake-executable planting not supported on Windows in this test") + } + // A trivial script that drains stdin and exits 0, using only shell builtins + // (the test empties PATH, so external commands like `cat` are unavailable). + script := "#!/bin/sh\nwhile read _; do :; done\nexit 0\n" + path := dir + string(os.PathSeparator) + name + if err := os.WriteFile(path, []byte(script), 0o755); err != nil { + t.Fatalf("write fake %s: %v", name, err) + } + t.Setenv("PATH", dir) + + if !Available() { + t.Fatal("Available() = false after planting fake utility") + } + if err := Copy("payload"); err != nil { + t.Errorf("Copy err = %v, want nil", err) + } +} diff --git a/pigo/internal/compaction/compact.go b/pigo/internal/compaction/compact.go new file mode 100644 index 0000000..935aa68 --- /dev/null +++ b/pigo/internal/compaction/compact.go @@ -0,0 +1,133 @@ +// This file (US-003) ties the compaction pieces together: given a message list +// and settings, it finds the cut point, extracts file operations from the +// summarized range, generates the structured summary, and returns a +// CompactionResult ready to be persisted as a session compaction entry and used +// to rebuild the agent context. Mirrors pi's prepareCompaction + compact. +package compaction + +import ( + "context" + "encoding/json" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// CompactionDetails records the files touched in the compacted history, stored +// alongside a compaction entry so a later iterative compaction can seed its +// file lists. Mirrors pi's CompactionDetails. +type CompactionDetails struct { + // ReadFiles are files read but not modified in the compacted range. + ReadFiles []string `json:"readFiles"` + // ModifiedFiles are files written or edited in the compacted range. + ModifiedFiles []string `json:"modifiedFiles"` +} + +// CompactionResult is the outcome of a compaction: the summary text (with file +// metadata appended), the index where retained history begins, the estimated +// tokens before compaction, and the extracted file details. Mirrors pi's +// CompactionResult, adapted to pigo's flat message list (index vs entry id). +type CompactionResult struct { + // Summary is the structured summary that replaces the compacted history. + Summary string + // FirstKeptIndex is the index of the first retained message. + FirstKeptIndex int + // TokensBefore is the estimated context tokens before compaction. + TokensBefore int + // Details holds the file operations extracted from the compacted range. + Details CompactionDetails +} + +// Compact prepares and generates a compaction over msgs. It cuts at +// FindCutPoint(msgs, settings.KeepRecentTokens), summarizes messages in +// [prevCompactionIndex+1, firstKeptIndex) (seeding file ops from a prior +// compaction's details when provided), and appends / +// metadata to the summary. previousSummary, when non-empty, switches to the +// iterative update template. +// +// prevCompactionIndex is the index of the last already-applied compaction point +// (-1 when none); summarization starts just after it so each compaction only +// covers the newly accumulated history. prevDetails carries that compaction's +// file lists so long-lived reads/edits survive across successive compactions. +// +// It returns (nil, nil) when there is nothing to compact (no valid cut point or +// an empty summarization range). +func Compact( + ctx context.Context, + stream provider.StreamFn, + model provider.Model, + msgs []agentcore.Message, + settings CompactionSettings, + prevCompactionIndex int, + prevDetails *CompactionDetails, + previousSummary string, + cfg provider.StreamConfig, +) (*CompactionResult, error) { + cut := FindCutPoint(msgs, settings.KeepRecentTokens) + + start := prevCompactionIndex + 1 + if start < 0 { + start = 0 + } + if start >= cut.FirstKeptIndex { + // Nothing new to summarize. + return nil, nil + } + toSummarize := msgs[start:cut.FirstKeptIndex] + + // Seed file ops from the previous compaction, then fold in this range. + ops := NewFileOps() + if prevDetails != nil { + for _, f := range prevDetails.ReadFiles { + ops.Read[f] = struct{}{} + } + for _, f := range prevDetails.ModifiedFiles { + ops.Edited[f] = struct{}{} + } + } + for _, m := range toSummarize { + extractFileOpsFromMessage(m, ops) + } + readFiles, modifiedFiles := computeFileLists(ops) + + summary, err := GenerateSummary(ctx, stream, model, toSummarize, settings.ReserveTokens, previousSummary, cfg) + if err != nil { + return nil, err + } + summary += formatFileOperations(readFiles, modifiedFiles) + + return &CompactionResult{ + Summary: summary, + FirstKeptIndex: cut.FirstKeptIndex, + TokensBefore: EstimateContextTokens(msgs).Tokens, + Details: CompactionDetails{ReadFiles: readFiles, ModifiedFiles: modifiedFiles}, + }, nil +} + +// Message builds the CompactionMessage to persist for this result: the summary +// text plus the estimated tokens-before and the file details (as raw JSON). +func (r *CompactionResult) Message(now int64) agentcore.CompactionMessage { + var details json.RawMessage + if b, err := json.Marshal(r.Details); err == nil { + details = b + } + return agentcore.CompactionMessage{ + RoleField: agentcore.RoleCompaction, + Summary: r.Summary, + TokensBefore: r.TokensBefore, + Details: details, + Timestamp: now, + } +} + +// RebuildContext returns the post-compaction message list: the compaction +// checkpoint followed by the retained recent messages (msgs[FirstKeptIndex:]). +// The summarized prefix is dropped and replaced by the single checkpoint, +// keeping context continuous while staying within the window. Mirrors pi's +// post-compact context reconstruction (summary entry + retained tail). +func (r *CompactionResult) RebuildContext(msgs []agentcore.Message, now int64) agentcore.MessageList { + out := make(agentcore.MessageList, 0, len(msgs)-r.FirstKeptIndex+1) + out = append(out, r.Message(now)) + out = append(out, msgs[r.FirstKeptIndex:]...) + return out +} diff --git a/pigo/internal/compaction/cutpoint.go b/pigo/internal/compaction/cutpoint.go new file mode 100644 index 0000000..4598180 --- /dev/null +++ b/pigo/internal/compaction/cutpoint.go @@ -0,0 +1,97 @@ +package compaction + +import "github.com/smallnest/pigo/internal/agentcore" + +// CutPointResult describes the cut selected for compaction, mirroring pi's +// CutPointResult (adapted to pigo's flat message list). +type CutPointResult struct { + // FirstKeptIndex is the index of the first message retained after + // compaction; everything before it is summarized. + FirstKeptIndex int + // TurnStartIndex is the index of the user message that starts the turn the + // cut falls inside, or -1 when the cut lands on a clean turn boundary. + TurnStartIndex int + // IsSplitTurn reports whether the cut splits an in-progress assistant turn. + IsSplitTurn bool +} + +// isValidCutPoint reports whether a message may serve as a cut point. A cut may +// land on a user or assistant message but never on a toolResult, because a +// toolResult must stay attached to the toolCall that produced it (pi semantics). +func isValidCutPoint(msg agentcore.Message) bool { + switch msg.Role() { + case agentcore.RoleUser, agentcore.RoleAssistant: + return true + default: // toolResult and any custom kinds are not cuttable. + return false + } +} + +// findValidCutPoints returns the indices of all messages that may serve as cut +// points, in ascending order. +func findValidCutPoints(msgs []agentcore.Message) []int { + var pts []int + for i, m := range msgs { + if isValidCutPoint(m) { + pts = append(pts, i) + } + } + return pts +} + +// findTurnStartIndex scans backwards from idx to find the user message that +// starts the turn containing idx, returning -1 if none is found. +func findTurnStartIndex(msgs []agentcore.Message, idx int) int { + for i := idx; i >= 0; i-- { + if msgs[i].Role() == agentcore.RoleUser { + return i + } + } + return -1 +} + +// FindCutPoint finds the compaction cut that keeps approximately +// keepRecentTokens worth of the most recent messages. It accumulates token +// estimates from the newest message backwards; once the retained budget is +// reached it snaps to the nearest valid cut point at or after that message, +// never splitting a toolCall from its toolResult. This mirrors pi's findCutPoint. +// +// When no valid cut point exists, it keeps everything (FirstKeptIndex 0). +func FindCutPoint(msgs []agentcore.Message, keepRecentTokens int) CutPointResult { + cutPoints := findValidCutPoints(msgs) + if len(cutPoints) == 0 { + return CutPointResult{FirstKeptIndex: 0, TurnStartIndex: -1, IsSplitTurn: false} + } + + // Default to the earliest valid cut point (keep as much as possible) when + // the retained budget is never reached. + cutIndex := cutPoints[0] + + accumulated := 0 + for i := len(msgs) - 1; i >= 0; i-- { + accumulated += EstimateTokens(msgs[i]) + if accumulated >= keepRecentTokens { + // Snap to the nearest valid cut point at or after i, so the kept + // window starts on a cuttable boundary. + for _, c := range cutPoints { + if c >= i { + cutIndex = c + break + } + } + break + } + } + + cutOnUser := msgs[cutIndex].Role() == agentcore.RoleUser + turnStart := -1 + if !cutOnUser { + turnStart = findTurnStartIndex(msgs, cutIndex) + } + + return CutPointResult{ + FirstKeptIndex: cutIndex, + TurnStartIndex: turnStart, + IsSplitTurn: !cutOnUser && turnStart != -1, + } +} diff --git a/pigo/internal/compaction/cutpoint_test.go b/pigo/internal/compaction/cutpoint_test.go new file mode 100644 index 0000000..3045378 --- /dev/null +++ b/pigo/internal/compaction/cutpoint_test.go @@ -0,0 +1,119 @@ +package compaction + +import ( + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// bigUser builds a user message whose estimate is ~tokens tokens (4 chars each). +func bigUser(tokens int) agentcore.UserMessage { + return userMsg(strings.Repeat("x", tokens*4)) +} + +func assistantToolCall(id, name string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{agentcore.NewToolCallContent(id, name, nil)}, + } +} + +func toolResult(id string) agentcore.ToolResultMessage { + return agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, + ToolCallID: id, + Content: agentcore.ContentList{agentcore.NewTextContent("result")}, + } +} + +func TestFindCutPointNoValidCutPoint(t *testing.T) { + // Only toolResult messages -> no valid cut point. + msgs := []agentcore.Message{toolResult("a"), toolResult("b")} + got := FindCutPoint(msgs, 100) + if got.FirstKeptIndex != 0 || got.TurnStartIndex != -1 || got.IsSplitTurn { + t.Fatalf("no-cutpoint: got %+v, want {0 -1 false}", got) + } +} + +func TestFindCutPointOnTurnBoundary(t *testing.T) { + // Three complete user/assistant turns, each user ~100 tokens. + // keepRecentTokens small so the cut lands on a clean user boundary. + msgs := []agentcore.Message{ + bigUser(100), // 0 + assistantMsg("a1", nil, ""), // 1 + bigUser(100), // 2 + assistantMsg("a2", nil, ""), // 3 + bigUser(100), // 4 + assistantMsg("a3", nil, ""), // 5 + } + // keepRecentTokens=50: walking back, msg[5] tiny, msg[4]=100 >= 50 at i=4. + // nearest cut point >= 4 is index 4 (a user message) -> clean boundary. + got := FindCutPoint(msgs, 50) + if got.FirstKeptIndex != 4 { + t.Fatalf("FirstKeptIndex: got %d, want 4", got.FirstKeptIndex) + } + if got.IsSplitTurn { + t.Fatalf("IsSplitTurn: got true, want false (landed on user boundary)") + } + if got.TurnStartIndex != -1 { + t.Fatalf("TurnStartIndex: got %d, want -1", got.TurnStartIndex) + } +} + +func TestFindCutPointSplitTurn(t *testing.T) { + // A turn starting with a user message, a big assistant reply, a toolCall and + // its toolResult. If the budget forces the cut onto the assistant message + // mid-turn, it's a split turn. + msgs := []agentcore.Message{ + userMsg("older turn"), // 0 + assistantMsg("older reply", nil, ""), // 1 + userMsg("current turn start"), // 2 user + assistantMsg(strings.Repeat("y", 400), nil, ""), // 3 assistant ~100 tokens + assistantToolCall("t1", "read"), // 4 assistant (valid cut) + toolResult("t1"), // 5 toolResult (not cuttable) + } + // keepRecentTokens=50: from end, msg[5]=~2, msg[4]~1, msg[3]=100 >= 50 at i=3. + // nearest cut point >= 3 is index 3 (assistant) -> split turn, turn start=2. + got := FindCutPoint(msgs, 50) + if got.FirstKeptIndex != 3 { + t.Fatalf("FirstKeptIndex: got %d, want 3", got.FirstKeptIndex) + } + if !got.IsSplitTurn { + t.Fatalf("IsSplitTurn: got false, want true (cut on assistant mid-turn)") + } + if got.TurnStartIndex != 2 { + t.Fatalf("TurnStartIndex: got %d, want 2", got.TurnStartIndex) + } +} + +func TestFindCutPointNeverCutsOnToolResult(t *testing.T) { + // Ensure a toolResult is never chosen even when it's the message where the + // budget is reached. + msgs := []agentcore.Message{ + userMsg("u0"), // 0 + assistantToolCall("t1", "grep"), // 1 + toolResult("t1"), // 2 (budget could land here) + assistantMsg("done", nil, ""), // 3 + } + got := FindCutPoint(msgs, 1) // tiny budget: reached at msg[3] + // cut must be a valid point (>=3 is index 3 assistant); never index 2. + if got.FirstKeptIndex == 2 { + t.Fatalf("cut landed on toolResult (index 2), which is illegal") + } + if msgs[got.FirstKeptIndex].Role() == agentcore.RoleToolResult { + t.Fatalf("FirstKeptIndex points at a toolResult: %d", got.FirstKeptIndex) + } +} + +func TestFindCutPointBudgetNeverReachedKeepsEarliest(t *testing.T) { + msgs := []agentcore.Message{ + userMsg("u0"), // 0 + assistantMsg("a0", nil, ""), // 1 + } + // Huge budget -> never reached -> keep from earliest valid cut point (0). + got := FindCutPoint(msgs, 1_000_000) + if got.FirstKeptIndex != 0 { + t.Fatalf("FirstKeptIndex: got %d, want 0 (earliest cut point)", got.FirstKeptIndex) + } +} diff --git a/pigo/internal/compaction/summary.go b/pigo/internal/compaction/summary.go new file mode 100644 index 0000000..ee3b1b2 --- /dev/null +++ b/pigo/internal/compaction/summary.go @@ -0,0 +1,367 @@ +// This file (US-003) covers summarization: the structured prompts, the +// conversation serializer, file-operation extraction, and GenerateSummary, +// which drives a provider stream to turn compacted history into a structured +// checkpoint summary. It mirrors pi's harness/compaction/compaction.ts prompts +// and utils.ts helpers, adapted to pigo's flat message list and Provider stream. +package compaction + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// SUMMARIZATION_SYSTEM_PROMPT instructs the model to only emit the structured +// summary and never continue the conversation. Verbatim from pi. +const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified. + +Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.` + +// summarizationPrompt is the first-time summary template (pi's SUMMARIZATION_PROMPT). +const summarizationPrompt = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work. + +Use this EXACT format: + +## Goal +[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.] + +## Constraints & Preferences +- [Any constraints, preferences, or requirements mentioned by user] +- [Or "(none)" if none were mentioned] + +## Progress +### Done +- [x] [Completed tasks/changes] + +### In Progress +- [ ] [Current work] + +### Blocked +- [Issues preventing progress, if any] + +## Key Decisions +- **[Decision]**: [Brief rationale] + +## Next Steps +1. [Ordered list of what should happen next] + +## Critical Context +- [Any data, examples, or references needed to continue] +- [Or "(none)" if not applicable] + +Keep each section concise. Preserve exact file paths, function names, and error messages.` + +// updateSummarizationPrompt incorporates new messages into an existing summary +// (pi's UPDATE_SUMMARIZATION_PROMPT). +const updateSummarizationPrompt = `The messages above are NEW conversation messages to incorporate into the existing summary provided in tags. + +Update the existing structured summary with new information. RULES: +- PRESERVE all existing information from the previous summary +- ADD new progress, decisions, and context from the new messages +- UPDATE the Progress section: move items from "In Progress" to "Done" when completed +- UPDATE "Next Steps" based on what was accomplished +- PRESERVE exact file paths, function names, and error messages +- If something is no longer relevant, you may remove it + +Use this EXACT format: + +## Goal +[Preserve existing goals, add new ones if the task expanded] + +## Constraints & Preferences +- [Preserve existing, add new ones discovered] + +## Progress +### Done +- [x] [Include previously done items AND newly completed items] + +### In Progress +- [ ] [Current work - update based on progress] + +### Blocked +- [Current blockers - remove if resolved] + +## Key Decisions +- **[Decision]**: [Brief rationale] (preserve all previous, add new) + +## Next Steps +1. [Update based on current state] + +## Critical Context +- [Preserve important context, add new if needed] + +Keep each section concise. Preserve exact file paths, function names, and error messages.` + +// FileOperations accumulates the file paths a compaction range touched, split +// by access kind. Mirrors pi's FileOperations. +type FileOperations struct { + // Read holds files read but not necessarily modified. + Read map[string]struct{} + // Written holds files written by full-file write operations. + Written map[string]struct{} + // Edited holds files modified by edit operations. + Edited map[string]struct{} +} + +// NewFileOps returns an empty file-operation accumulator. +func NewFileOps() FileOperations { + return FileOperations{ + Read: map[string]struct{}{}, + Written: map[string]struct{}{}, + Edited: map[string]struct{}{}, + } +} + +// extractFileOpsFromMessage adds file operations from an assistant message's +// tool calls to the accumulator, keyed on the tool name (read/write/edit) and +// the string "path" argument. Non-assistant messages are ignored, matching pi. +func extractFileOpsFromMessage(msg agentcore.Message, ops FileOperations) { + a, ok := msg.(agentcore.AssistantMessage) + if !ok { + return + } + for _, call := range a.ToolCalls() { + path := toolCallPath(call.Arguments) + if path == "" { + continue + } + switch call.Name { + case "read": + ops.Read[path] = struct{}{} + case "write": + ops.Written[path] = struct{}{} + case "edit": + ops.Edited[path] = struct{}{} + } + } +} + +// toolCallPath extracts a string "path" argument from a tool call's raw +// arguments, returning "" when absent or not a string. +func toolCallPath(args json.RawMessage) string { + if len(args) == 0 { + return "" + } + var decoded struct { + Path string `json:"path"` + } + if err := json.Unmarshal(args, &decoded); err != nil { + return "" + } + return decoded.Path +} + +// computeFileLists derives sorted read-only and modified (edited∪written) file +// lists, excluding modified files from the read-only list. Mirrors pi. +func computeFileLists(ops FileOperations) (readFiles, modifiedFiles []string) { + modified := map[string]struct{}{} + for f := range ops.Edited { + modified[f] = struct{}{} + } + for f := range ops.Written { + modified[f] = struct{}{} + } + for f := range ops.Read { + if _, isMod := modified[f]; !isMod { + readFiles = append(readFiles, f) + } + } + for f := range modified { + modifiedFiles = append(modifiedFiles, f) + } + sort.Strings(readFiles) + sort.Strings(modifiedFiles) + return readFiles, modifiedFiles +} + +// formatFileOperations renders the file lists as / +// metadata blocks appended to a summary, or "" when both lists are empty. +func formatFileOperations(readFiles, modifiedFiles []string) string { + var sections []string + if len(readFiles) > 0 { + sections = append(sections, "\n"+strings.Join(readFiles, "\n")+"\n") + } + if len(modifiedFiles) > 0 { + sections = append(sections, "\n"+strings.Join(modifiedFiles, "\n")+"\n") + } + if len(sections) == 0 { + return "" + } + return "\n\n" + strings.Join(sections, "\n\n") +} + +// toolResultMaxChars caps a tool result's serialized text in the summarization +// prompt, matching pi's TOOL_RESULT_MAX_CHARS. +const toolResultMaxChars = 2000 + +// truncateForSummary caps text at maxChars, appending a truncation marker. +func truncateForSummary(text string, maxChars int) string { + if len(text) <= maxChars { + return text + } + return fmt.Sprintf("%s\n\n[... %d more characters truncated]", text[:maxChars], len(text)-maxChars) +} + +// textOf concatenates the text blocks of a content list. +func textOf(content agentcore.ContentList) string { + var b strings.Builder + for _, c := range content { + if t, ok := c.(agentcore.TextContent); ok { + b.WriteString(t.Text) + } + } + return b.String() +} + +// serializeConversation renders messages as a plain-text transcript for the +// summarization prompt, mirroring pi's serializeConversation: user text, +// assistant thinking/text/tool-call lines, and truncated tool results. +func serializeConversation(msgs []agentcore.Message) string { + var parts []string + for _, msg := range msgs { + switch m := msg.(type) { + case agentcore.UserMessage: + if s := textOf(m.Content); s != "" { + parts = append(parts, "[User]: "+s) + } + case agentcore.AssistantMessage: + var textParts, thinkingParts, toolCalls []string + for _, block := range m.Content { + switch c := block.(type) { + case agentcore.TextContent: + textParts = append(textParts, c.Text) + case agentcore.ThinkingContent: + thinkingParts = append(thinkingParts, c.Thinking) + case agentcore.ToolCallContent: + toolCalls = append(toolCalls, formatToolCall(c)) + } + } + if len(thinkingParts) > 0 { + parts = append(parts, "[Assistant thinking]: "+strings.Join(thinkingParts, "\n")) + } + if len(textParts) > 0 { + parts = append(parts, "[Assistant]: "+strings.Join(textParts, "\n")) + } + if len(toolCalls) > 0 { + parts = append(parts, "[Assistant tool calls]: "+strings.Join(toolCalls, "; ")) + } + case agentcore.ToolResultMessage: + if s := textOf(m.Content); s != "" { + parts = append(parts, "[Tool result]: "+truncateForSummary(s, toolResultMaxChars)) + } + } + } + return strings.Join(parts, "\n\n") +} + +// formatToolCall renders a tool call as name(key=value, ...) with each value +// JSON-encoded, mirroring pi's serialization of tool-call arguments. +func formatToolCall(c agentcore.ToolCallContent) string { + if len(c.Arguments) == 0 { + return c.Name + "()" + } + var m map[string]json.RawMessage + if err := json.Unmarshal(c.Arguments, &m); err != nil { + // Not an object: fall back to the raw argument text. + return fmt.Sprintf("%s(%s)", c.Name, string(c.Arguments)) + } + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) // deterministic order (Go map iteration is random) + pairs := make([]string, 0, len(keys)) + for _, k := range keys { + pairs = append(pairs, k+"="+string(m[k])) + } + return fmt.Sprintf("%s(%s)", c.Name, strings.Join(pairs, ", ")) +} + +// GenerateSummary drives a provider stream to summarize msgs into a structured +// checkpoint. When previousSummary is non-empty it uses the update template and +// embeds the prior summary in tags; otherwise it uses the +// first-time template. maxTokens is bounded to min(0.8*reserveTokens, +// model.MaxOutputTokens) as in pi. The returned text is the concatenation of +// the assistant response's text blocks. A terminal error/aborted response is +// surfaced as an error. +func GenerateSummary( + ctx context.Context, + stream provider.StreamFn, + model provider.Model, + msgs []agentcore.Message, + reserveTokens int, + previousSummary string, + cfg provider.StreamConfig, +) (string, error) { + base := summarizationPrompt + if previousSummary != "" { + base = updateSummarizationPrompt + } + + conversation := serializeConversation(msgs) + var b strings.Builder + b.WriteString("\n") + b.WriteString(conversation) + b.WriteString("\n\n\n") + if previousSummary != "" { + b.WriteString("\n") + b.WriteString(previousSummary) + b.WriteString("\n\n\n") + } + b.WriteString(base) + + promptMsg := agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(b.String())}, + } + + // Bound the summary output to min(0.8*reserveTokens, model max output). + maxTokens := (reserveTokens * 8) / 10 + if model.MaxOutputTokens > 0 && model.MaxOutputTokens < maxTokens { + maxTokens = model.MaxOutputTokens + } + extra := map[string]any{} + for k, v := range cfg.Extra { + extra[k] = v + } + if maxTokens > 0 { + extra["max_tokens"] = maxTokens + } + cfg.Extra = extra + + llm := provider.LlmContext{ + SystemPrompt: SUMMARIZATION_SYSTEM_PROMPT, + Messages: agentcore.MessageList{promptMsg}, + } + s, err := stream(ctx, model.ID, llm, cfg) + if err != nil { + return "", fmt.Errorf("compaction: build summary stream: %w", err) + } + + // Drain events so the stream's result is populated, then read the final + // message. Failures ride the stream as a terminal message per the provider + // contract, so inspect StopReason rather than only the returned error. + for range s.Events() { + } + final, resErr := s.Result(ctx) + if resErr != nil { + return "", fmt.Errorf("compaction: summary stream: %w", resErr) + } + switch final.StopReason { + case agentcore.StopReasonAborted: + return "", fmt.Errorf("compaction: summarization aborted: %s", final.ErrorMessage) + case agentcore.StopReasonError: + return "", fmt.Errorf("compaction: summarization failed: %s", final.ErrorMessage) + } + + summary := textOf(final.Content) + if strings.TrimSpace(summary) == "" { + return "", fmt.Errorf("compaction: summarization produced empty output") + } + return summary, nil +} diff --git a/pigo/internal/compaction/summary_test.go b/pigo/internal/compaction/summary_test.go new file mode 100644 index 0000000..dc3c6b7 --- /dev/null +++ b/pigo/internal/compaction/summary_test.go @@ -0,0 +1,274 @@ +package compaction + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +func TestExtractFileOpsAndLists(t *testing.T) { + readArgs, _ := json.Marshal(map[string]string{"path": "a.go"}) + writeArgs, _ := json.Marshal(map[string]string{"path": "b.go"}) + editArgs, _ := json.Marshal(map[string]string{"path": "a.go"}) // a.go also edited -> modified wins + msgs := []agentcore.Message{ + assistantToolCall("1", "read"), + assistantToolCall("2", "write"), + assistantToolCall("3", "edit"), + } + // Attach args by rebuilding with arguments. + msgs[0] = agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewToolCallContent("1", "read", readArgs)}} + msgs[1] = agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewToolCallContent("2", "write", writeArgs)}} + msgs[2] = agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewToolCallContent("3", "edit", editArgs)}} + + ops := NewFileOps() + for _, m := range msgs { + extractFileOpsFromMessage(m, ops) + } + read, modified := computeFileLists(ops) + // a.go was read AND edited -> only in modified; b.go written -> modified. + if len(read) != 0 { + t.Fatalf("readFiles: got %v, want []", read) + } + if strings.Join(modified, ",") != "a.go,b.go" { + t.Fatalf("modifiedFiles: got %v, want [a.go b.go]", modified) + } +} + +func TestFormatFileOperations(t *testing.T) { + if got := formatFileOperations(nil, nil); got != "" { + t.Fatalf("empty: got %q, want empty", got) + } + got := formatFileOperations([]string{"r.go"}, []string{"m.go"}) + if !strings.Contains(got, "\nr.go\n") { + t.Fatalf("missing read-files block: %q", got) + } + if !strings.Contains(got, "\nm.go\n") { + t.Fatalf("missing modified-files block: %q", got) + } +} + +func TestSerializeConversation(t *testing.T) { + args, _ := json.Marshal(map[string]any{"path": "x.go", "n": 1}) + msgs := []agentcore.Message{ + userMsg("hello"), + agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{ + agentcore.NewThinkingContent("thinking hard"), + agentcore.NewTextContent("here goes"), + agentcore.NewToolCallContent("t1", "read", args), + }, + }, + toolResult("t1"), + } + got := serializeConversation(msgs) + for _, want := range []string{ + "[User]: hello", + "[Assistant thinking]: thinking hard", + "[Assistant]: here goes", + "[Assistant tool calls]: read(", + "[Tool result]: result", + } { + if !strings.Contains(got, want) { + t.Fatalf("serialize missing %q in:\n%s", want, got) + } + } +} + +func TestTruncateForSummary(t *testing.T) { + if got := truncateForSummary("short", 100); got != "short" { + t.Fatalf("no truncation expected: %q", got) + } + long := strings.Repeat("z", 2500) + got := truncateForSummary(long, toolResultMaxChars) + if !strings.Contains(got, "more characters truncated") { + t.Fatalf("expected truncation marker: %q", got[len(got)-60:]) + } +} + +// fakeStreamFn returns a StreamFn that yields a single done event with the +// given assistant message, capturing the LlmContext it was called with. +func fakeStreamFn(final agentcore.AssistantMessage, capture *provider.LlmContext) provider.StreamFn { + return func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + if capture != nil { + *capture = llm + } + s := provider.NewAssistantMessageEventStream(4) + go func() { + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: final}) + s.SetResult(final) + s.Close() + }() + return s, nil + } +} + +func assistantText(text string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + StopReason: agentcore.StopReasonEndTurn, + } +} + +func TestGenerateSummaryFirstTime(t *testing.T) { + var captured provider.LlmContext + stream := fakeStreamFn(assistantText("## Goal\ndo the thing"), &captured) + model := provider.Model{ID: "m", MaxOutputTokens: 8000} + msgs := []agentcore.Message{userMsg("please do X")} + + got, err := GenerateSummary(context.Background(), stream, model, msgs, 16384, "", provider.StreamConfig{}) + if err != nil { + t.Fatalf("GenerateSummary: %v", err) + } + if !strings.Contains(got, "## Goal") { + t.Fatalf("summary text: %q", got) + } + // System prompt must be the summarization system prompt. + if captured.SystemPrompt != SUMMARIZATION_SYSTEM_PROMPT { + t.Fatalf("system prompt mismatch") + } + // First-time prompt uses the non-update template and wraps the conversation. + promptText := textOf(captured.Messages[0].(agentcore.UserMessage).Content) + if !strings.Contains(promptText, "") || strings.Contains(promptText, "") { + t.Fatalf("first-time prompt shape wrong:\n%s", promptText) + } + if !strings.Contains(promptText, "Create a structured context checkpoint") { + t.Fatalf("expected first-time template") + } +} + +func TestGenerateSummaryUpdateUsesPrevious(t *testing.T) { + var captured provider.LlmContext + stream := fakeStreamFn(assistantText("updated summary"), &captured) + model := provider.Model{ID: "m"} + msgs := []agentcore.Message{userMsg("more work")} + + _, err := GenerateSummary(context.Background(), stream, model, msgs, 16384, "PRIOR SUMMARY", provider.StreamConfig{}) + if err != nil { + t.Fatalf("GenerateSummary: %v", err) + } + promptText := textOf(captured.Messages[0].(agentcore.UserMessage).Content) + if !strings.Contains(promptText, "\nPRIOR SUMMARY") { + t.Fatalf("update prompt should embed previous summary:\n%s", promptText) + } + if !strings.Contains(promptText, "NEW conversation messages to incorporate") { + t.Fatalf("expected update template") + } +} + +func TestGenerateSummaryMaxTokensCap(t *testing.T) { + var gotMax int + stream := func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + if v, ok := cfg.Extra["max_tokens"].(int); ok { + gotMax = v + } + s := provider.NewAssistantMessageEventStream(2) + go func() { + m := assistantText("ok") + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: m}) + s.SetResult(m) + s.Close() + }() + return s, nil + } + // 0.8 * 16384 = 13107, but model max output is 5000 -> cap at 5000. + model := provider.Model{ID: "m", MaxOutputTokens: 5000} + _, err := GenerateSummary(context.Background(), stream, model, []agentcore.Message{userMsg("x")}, 16384, "", provider.StreamConfig{}) + if err != nil { + t.Fatalf("GenerateSummary: %v", err) + } + if gotMax != 5000 { + t.Fatalf("max_tokens: got %d, want 5000", gotMax) + } +} + +func TestGenerateSummaryErrorStopReason(t *testing.T) { + errMsg := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonError, ErrorMessage: "boom"} + stream := fakeStreamFn(errMsg, nil) + _, err := GenerateSummary(context.Background(), stream, provider.Model{ID: "m"}, []agentcore.Message{userMsg("x")}, 16384, "", provider.StreamConfig{}) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("expected error containing 'boom', got %v", err) + } +} + +func TestCompactRebuildsContext(t *testing.T) { + readArgs, _ := json.Marshal(map[string]string{"path": "old.go"}) + msgs := []agentcore.Message{ + userMsg("turn one"), // 0 + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewToolCallContent("t1", "read", readArgs)}}, // 1 + toolResult("t1"), // 2 + bigUser(100), // 3 + assistantMsg("recent", nil, ""), // 4 + } + stream := fakeStreamFn(assistantText("## Goal\nx"), nil) + // Small keepRecentTokens so the cut lands on the recent bigUser(100) turn, + // leaving the earlier read/toolResult prefix to be summarized. + settings := CompactionSettings{Enabled: true, ReserveTokens: 16384, KeepRecentTokens: 50} + res, err := Compact(context.Background(), stream, provider.Model{ID: "m"}, msgs, settings, -1, nil, "", provider.StreamConfig{}) + if err != nil { + t.Fatalf("Compact: %v", err) + } + if res == nil { + t.Fatal("Compact returned nil result") + } + // old.go was read in the summarized prefix. + if strings.Join(res.Details.ReadFiles, ",") != "old.go" { + t.Fatalf("readFiles: got %v, want [old.go]", res.Details.ReadFiles) + } + if !strings.Contains(res.Summary, "") { + t.Fatalf("summary should carry file metadata: %q", res.Summary) + } + rebuilt := res.RebuildContext(msgs, 123) + if rebuilt[0].Role() != agentcore.RoleCompaction { + t.Fatalf("first rebuilt message must be compaction, got %s", rebuilt[0].Role()) + } + // Retained tail begins at FirstKeptIndex. + if len(rebuilt) != 1+(len(msgs)-res.FirstKeptIndex) { + t.Fatalf("rebuilt length: got %d", len(rebuilt)) + } +} + +func TestCompactNothingToSummarize(t *testing.T) { + // prevCompactionIndex already at/after the cut -> nil result. + msgs := []agentcore.Message{userMsg("a"), assistantMsg("b", nil, "")} + stream := fakeStreamFn(assistantText("unused"), nil) + res, err := Compact(context.Background(), stream, provider.Model{ID: "m"}, msgs, DefaultCompactionSettings, 5, nil, "", provider.StreamConfig{}) + if err != nil { + t.Fatalf("Compact: %v", err) + } + if res != nil { + t.Fatalf("expected nil result when nothing to summarize, got %+v", res) + } +} + +func TestCompactionMessageRoundTrip(t *testing.T) { + details, _ := json.Marshal(CompactionDetails{ReadFiles: []string{"a"}, ModifiedFiles: []string{"b"}}) + cm := agentcore.CompactionMessage{ + RoleField: agentcore.RoleCompaction, + Summary: "the summary", + TokensBefore: 42, + Details: details, + Timestamp: 7, + } + list := agentcore.MessageList{cm} + raw, err := json.Marshal(list) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var back agentcore.MessageList + if err := json.Unmarshal(raw, &back); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if len(back) != 1 || back[0].Role() != agentcore.RoleCompaction { + t.Fatalf("round-trip role: %+v", back) + } + got := back[0].(agentcore.CompactionMessage) + if got.Summary != "the summary" || got.TokensBefore != 42 { + t.Fatalf("round-trip fields: %+v", got) + } +} diff --git a/pigo/internal/compaction/tokens.go b/pigo/internal/compaction/tokens.go new file mode 100644 index 0000000..8297524 --- /dev/null +++ b/pigo/internal/compaction/tokens.go @@ -0,0 +1,189 @@ +// Package compaction implements context-window token accounting and the +// decision of when a long session must be compacted, mirroring pi's +// harness/compaction/compaction.ts. +// +// This file (US-001) covers the token side: estimating a message's token +// footprint from a character heuristic, deriving the current context-token +// usage (preferring provider-reported Usage over estimation), and the +// ShouldCompact threshold check. Cut-point finding (US-002) and summarization +// (US-003) live in sibling files. +package compaction + +import ( + "encoding/json" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// CompactionSettings holds the thresholds and retention knobs for compaction, +// mirroring pi's CompactionSettings. +type CompactionSettings struct { + // Enabled gates automatic compaction decisions. + Enabled bool + // ReserveTokens is reserved for the summary prompt and its output; the + // effective usable window is contextWindow - ReserveTokens. + ReserveTokens int + // KeepRecentTokens is the approximate recent-context token budget to retain + // after compaction (consumed by FindCutPoint in US-002). + KeepRecentTokens int +} + +// DefaultCompactionSettings matches pi's DEFAULT_COMPACTION_SETTINGS. +var DefaultCompactionSettings = CompactionSettings{ + Enabled: true, + ReserveTokens: 16384, + KeepRecentTokens: 20000, +} + +// estimatedImageChars is the fixed character budget attributed to an image +// block, matching pi's ESTIMATED_IMAGE_CHARS. +const estimatedImageChars = 4800 + +// charsPerToken is the conservative characters-per-token divisor pi uses. +const charsPerToken = 4 + +// ceilDiv returns ceil(a / b) for non-negative a and positive b. +func ceilDiv(a, b int) int { + if a <= 0 { + return 0 + } + return (a + b - 1) / b +} + +// contentListChars sums the character footprint of a content list, counting +// text/thinking/toolCall blocks by their text length and each image block as a +// fixed estimatedImageChars, mirroring pi's estimateTextAndImageContentChars +// plus its assistant-block handling. +func contentListChars(content agentcore.ContentList) int { + chars := 0 + for _, block := range content { + switch c := block.(type) { + case agentcore.TextContent: + chars += len(c.Text) + case agentcore.ThinkingContent: + chars += len(c.Thinking) + case agentcore.ToolCallContent: + // name + serialized arguments, matching pi's toolCall accounting. + chars += len(c.Name) + if len(c.Arguments) > 0 { + chars += len(c.Arguments) + } else { + // nil/empty RawMessage serializes to "null" downstream. + b, _ := json.Marshal(json.RawMessage(c.Arguments)) + chars += len(b) + } + case agentcore.ImageContent: + chars += estimatedImageChars + } + } + return chars +} + +// EstimateTokens returns a conservative token estimate for one message using +// the same character heuristic as pi's estimateTokens (ceil(chars / 4)). +func EstimateTokens(msg agentcore.Message) int { + switch m := msg.(type) { + case agentcore.UserMessage: + return ceilDiv(contentListChars(m.Content), charsPerToken) + case agentcore.AssistantMessage: + return ceilDiv(contentListChars(m.Content), charsPerToken) + case agentcore.ToolResultMessage: + return ceilDiv(contentListChars(m.Content), charsPerToken) + case agentcore.CompactionMessage: + // A compaction checkpoint replays as its summary text; estimate from it. + return ceilDiv(len(m.Summary), charsPerToken) + default: + return 0 + } +} + +// calculateContextTokens derives total context tokens from a provider usage +// block. pigo's Usage only reports input/output, so we sum them (pi additionally +// folds cache read/write, which pigo does not track). +func calculateContextTokens(u agentcore.Usage) int { + return u.InputTokens + u.OutputTokens +} + +// assistantUsage returns a usable Usage from an assistant message, skipping +// aborted/error responses and zero-token usage, mirroring pi's getAssistantUsage. +func assistantUsage(msg agentcore.Message) (agentcore.Usage, bool) { + a, ok := msg.(agentcore.AssistantMessage) + if !ok || a.Usage == nil { + return agentcore.Usage{}, false + } + if a.StopReason == agentcore.StopReasonAborted || a.StopReason == agentcore.StopReasonError { + return agentcore.Usage{}, false + } + if calculateContextTokens(*a.Usage) <= 0 { + return agentcore.Usage{}, false + } + return *a.Usage, true +} + +// ContextUsageEstimate reports the derived context-token usage for a message +// list, mirroring pi's ContextUsageEstimate. +type ContextUsageEstimate struct { + // Tokens is the estimated total context tokens. + Tokens int + // UsageTokens is the tokens reported by the most recent assistant usage block. + UsageTokens int + // TrailingTokens is the estimated tokens after that usage block. + TrailingTokens int + // LastUsageIndex is the index of the message that provided usage, or -1 when + // none exists. + LastUsageIndex int +} + +// EstimateContextTokens computes context-token usage for messages, preferring +// the most recent valid assistant Usage block and estimating only the messages +// that follow it. When no usage is available it estimates every message. This +// mirrors pi's estimateContextTokens. +func EstimateContextTokens(msgs []agentcore.Message) ContextUsageEstimate { + lastIdx := -1 + var lastUsage agentcore.Usage + for i := len(msgs) - 1; i >= 0; i-- { + if u, ok := assistantUsage(msgs[i]); ok { + lastIdx = i + lastUsage = u + break + } + } + + if lastIdx < 0 { + estimated := 0 + for _, m := range msgs { + estimated += EstimateTokens(m) + } + return ContextUsageEstimate{ + Tokens: estimated, + UsageTokens: 0, + TrailingTokens: estimated, + LastUsageIndex: -1, + } + } + + usageTokens := calculateContextTokens(lastUsage) + trailing := 0 + for i := lastIdx + 1; i < len(msgs); i++ { + trailing += EstimateTokens(msgs[i]) + } + return ContextUsageEstimate{ + Tokens: usageTokens + trailing, + UsageTokens: usageTokens, + TrailingTokens: trailing, + LastUsageIndex: lastIdx, + } +} + +// ShouldCompact reports whether context usage has exceeded the usable window, +// matching pi: contextTokens > contextWindow - reserveTokens. Disabled settings +// or a non-positive contextWindow (unknown) never trigger compaction. +func ShouldCompact(contextTokens, contextWindow int, settings CompactionSettings) bool { + if !settings.Enabled { + return false + } + if contextWindow <= 0 { + return false + } + return contextTokens > contextWindow-settings.ReserveTokens +} diff --git a/pigo/internal/compaction/tokens_test.go b/pigo/internal/compaction/tokens_test.go new file mode 100644 index 0000000..78e7d08 --- /dev/null +++ b/pigo/internal/compaction/tokens_test.go @@ -0,0 +1,189 @@ +package compaction + +import ( + "encoding/json" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +func userMsg(text string) agentcore.UserMessage { + return agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + } +} + +func assistantMsg(text string, usage *agentcore.Usage, stop string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + Usage: usage, + StopReason: stop, + } +} + +func TestEstimateTokensText(t *testing.T) { + // 8 chars / 4 = 2 tokens. + if got := EstimateTokens(userMsg("abcdefgh")); got != 2 { + t.Fatalf("text estimate: got %d, want 2", got) + } + // 9 chars -> ceil(9/4) = 3. + if got := EstimateTokens(userMsg("abcdefghi")); got != 3 { + t.Fatalf("ceil estimate: got %d, want 3", got) + } + // empty -> 0. + if got := EstimateTokens(userMsg("")); got != 0 { + t.Fatalf("empty estimate: got %d, want 0", got) + } +} + +func TestEstimateTokensImage(t *testing.T) { + m := agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewImageContent("data", "image/png")}, + } + // estimatedImageChars (4800) / 4 = 1200. + if got := EstimateTokens(m); got != 1200 { + t.Fatalf("image estimate: got %d, want 1200", got) + } +} + +func TestEstimateTokensAssistantToolCall(t *testing.T) { + args := json.RawMessage(`{"path":"a.go"}`) // 15 chars + m := agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{ + agentcore.NewToolCallContent("id1", "read", args), // name "read" = 4 chars + }, + } + // (4 + 15) / 4 = ceil(19/4) = 5. + if got := EstimateTokens(m); got != 5 { + t.Fatalf("toolcall estimate: got %d, want 5", got) + } +} + +func TestEstimateTokensUnknownRoleZero(t *testing.T) { + if got := EstimateTokens(agentcore.ToolResultMessage{}); got != 0 { + t.Fatalf("empty tool result: got %d, want 0", got) + } +} + +func TestEstimateContextTokensNoUsageFallsBackToEstimate(t *testing.T) { + msgs := []agentcore.Message{ + userMsg("abcdefgh"), // 2 + assistantMsg("abcd", nil, "end_turn"), // 1 + } + est := EstimateContextTokens(msgs) + if est.LastUsageIndex != -1 { + t.Fatalf("LastUsageIndex: got %d, want -1", est.LastUsageIndex) + } + if est.Tokens != 3 { + t.Fatalf("Tokens: got %d, want 3", est.Tokens) + } + if est.TrailingTokens != 3 || est.UsageTokens != 0 { + t.Fatalf("trailing/usage: got %d/%d, want 3/0", est.TrailingTokens, est.UsageTokens) + } +} + +func TestEstimateContextTokensPrefersUsage(t *testing.T) { + usage := &agentcore.Usage{InputTokens: 1000, OutputTokens: 200} + msgs := []agentcore.Message{ + userMsg("first"), + assistantMsg("reply", usage, "end_turn"), + userMsg("abcdefgh"), // trailing: 2 tokens + } + est := EstimateContextTokens(msgs) + if est.LastUsageIndex != 1 { + t.Fatalf("LastUsageIndex: got %d, want 1", est.LastUsageIndex) + } + if est.UsageTokens != 1200 { + t.Fatalf("UsageTokens: got %d, want 1200", est.UsageTokens) + } + if est.TrailingTokens != 2 { + t.Fatalf("TrailingTokens: got %d, want 2", est.TrailingTokens) + } + if est.Tokens != 1202 { + t.Fatalf("Tokens: got %d, want 1202", est.Tokens) + } +} + +func TestEstimateContextTokensSkipsAbortedAndErrorUsage(t *testing.T) { + good := &agentcore.Usage{InputTokens: 500, OutputTokens: 0} + bad := &agentcore.Usage{InputTokens: 9999, OutputTokens: 0} + msgs := []agentcore.Message{ + assistantMsg("ok", good, "end_turn"), + assistantMsg("aborted", bad, agentcore.StopReasonAborted), + assistantMsg("errored", bad, agentcore.StopReasonError), + } + est := EstimateContextTokens(msgs) + if est.LastUsageIndex != 0 { + t.Fatalf("LastUsageIndex: got %d, want 0 (should skip aborted/error)", est.LastUsageIndex) + } + if est.UsageTokens != 500 { + t.Fatalf("UsageTokens: got %d, want 500", est.UsageTokens) + } +} + +func TestEstimateContextTokensSkipsZeroUsage(t *testing.T) { + zero := &agentcore.Usage{InputTokens: 0, OutputTokens: 0} + msgs := []agentcore.Message{ + userMsg("abcdefgh"), // 2 + assistantMsg("reply", zero, "end_turn"), + } + est := EstimateContextTokens(msgs) + // zero usage is ignored, so falls back to full estimation. + if est.LastUsageIndex != -1 { + t.Fatalf("LastUsageIndex: got %d, want -1", est.LastUsageIndex) + } +} + +func TestShouldCompactThresholdBoundaries(t *testing.T) { + s := CompactionSettings{Enabled: true, ReserveTokens: 16384} + window := 200000 + usable := window - s.ReserveTokens // 183616 + + tests := []struct { + name string + contextTokens int + want bool + }{ + {"far below", 1000, false}, + {"equal to usable", usable, false}, // strictly greater required + {"one over usable", usable + 1, true}, + {"far over", window * 2, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := ShouldCompact(tt.contextTokens, window, s); got != tt.want { + t.Fatalf("ShouldCompact(%d): got %v, want %v", tt.contextTokens, got, tt.want) + } + }) + } +} + +func TestShouldCompactDisabled(t *testing.T) { + s := CompactionSettings{Enabled: false, ReserveTokens: 16384} + if ShouldCompact(1_000_000, 200000, s) { + t.Fatal("disabled settings must never compact") + } +} + +func TestShouldCompactUnknownWindow(t *testing.T) { + s := CompactionSettings{Enabled: true, ReserveTokens: 16384} + if ShouldCompact(1_000_000, 0, s) { + t.Fatal("unknown (0) context window must never compact") + } +} + +func TestDefaultCompactionSettings(t *testing.T) { + if DefaultCompactionSettings.ReserveTokens != 16384 { + t.Fatalf("ReserveTokens: got %d, want 16384", DefaultCompactionSettings.ReserveTokens) + } + if DefaultCompactionSettings.KeepRecentTokens != 20000 { + t.Fatalf("KeepRecentTokens: got %d, want 20000", DefaultCompactionSettings.KeepRecentTokens) + } + if !DefaultCompactionSettings.Enabled { + t.Fatal("default settings should be enabled") + } +} diff --git a/pigo/internal/dream/apply.go b/pigo/internal/dream/apply.go new file mode 100644 index 0000000..8404a61 --- /dev/null +++ b/pigo/internal/dream/apply.go @@ -0,0 +1,247 @@ +package dream + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// This file wires the llmConsolidator to a real provider (the main-session +// model, SPEC Q3) and holds the MEMORY.md index cleanup the Runner runs after +// writeback. The provider plumbing mirrors internal/cli/run.SetupEnv: resolve a +// Provider for the model/base-url/protocol/provider tuple, resolve the API key +// through a CredentialStore (--api-key override → env → config), and run a +// single StreamCompletion, draining the event stream to the final text. + +// NewLLMConsolidator builds the production Consolidator backed by the given +// model configuration — the same tuple cmd/pigo resolves for the main session +// (CLI flags overlaid with config.toml). It resolves the Provider once so every +// Consolidate call reuses it. A resolution failure (bad model / missing +// provider) is returned so the caller can decide whether to fall back to the +// no-op Consolidator or fail the run. +func NewLLMConsolidator(model, baseURL, protocol, providerName, apiKey string, thinking agentcore.ThinkingLevel) (Consolidator, error) { + complete, err := newModelCompleter(model, baseURL, protocol, providerName, apiKey, thinking) + if err != nil { + return nil, err + } + return &llmConsolidator{complete: complete}, nil +} + +// newModelCompleter resolves the provider and returns a completeFn that performs +// one non-streaming-consuming completion: it sends the system+user prompt as a +// single user turn (no tools — the dream agent only reasons and replies) and +// returns the concatenated assistant text. A hard "cannot build the stream" +// error is returned directly; a runtime failure rides the stream as a terminal +// error event whose message we convert to an error (so the Runner marks the run +// failed rather than silently deleting nothing, SPEC §5.5). +func newModelCompleter(model, baseURL, protocol, providerName, apiKey string, thinking agentcore.ThinkingLevel) (completeFn, error) { + prov, resolvedName, err := provider.ResolveProvider(model, baseURL, protocol, providerName, os.Getenv) + if err != nil { + return nil, fmt.Errorf("dream: resolve provider: %w", err) + } + creds := provider.NewCredentialStore(nil) + creds.SetOverride(resolvedName, apiKey) + + return func(ctx context.Context, systemPrompt, userPrompt string) (string, error) { + key := creds.GetAPIKey(ctx, resolvedName) + llm := provider.LlmContext{ + SystemPrompt: systemPrompt, + Messages: agentcore.MessageList{ + agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(userPrompt)}, + }, + }, + } + stream, err := prov.StreamCompletion(ctx, provider.CompletionRequest{ + Model: model, + Context: llm, + Config: provider.StreamConfig{APIKey: key, ThinkingLevel: thinking}, + }) + if err != nil { + return "", err + } + final, err := drainToMessage(ctx, stream) + if err != nil { + return "", err + } + if final.StopReason == agentcore.StopReasonError { + if final.ErrorMessage != "" { + return "", fmt.Errorf("model error: %s", final.ErrorMessage) + } + return "", fmt.Errorf("model returned an error response") + } + return agentcore.ContentToText(final.Content), nil + }, nil +} + +// drainToMessage consumes the provider event stream to completion and returns +// the terminal assistant message. It mirrors the loop's stream-drain contract: +// the done/error event carries the final message; if the stream closes without +// one, it falls back to the stream Result. Draining is required because the +// producer blocks on the event channel until consumed. +func drainToMessage(ctx context.Context, stream *provider.AssistantMessageEventStream) (agentcore.AssistantMessage, error) { + for ev := range stream.Events() { + switch e := ev.(type) { + case provider.StreamDoneEvent: + return e.Message, nil + case provider.StreamErrorEvent: + return e.Message, nil + } + } + final, err := stream.Result(ctx) + if err != nil { + return agentcore.AssistantMessage{}, err + } + return final, nil +} + +// updateScopeIndexes rewrites each affected scope's MEMORY.md to drop any line +// that references a now-deleted memory file, keeping the index consistent with +// the entries on disk and free of dangling links (PRD US-003). It is safe to +// call when no MEMORY.md exists (no-op) and when deleted is empty. Each rewrite +// is atomic (temp+rename) and guarded by withinScope, so it cannot escape the +// memory store. +func updateScopeIndexes(memoryRoot, projectDir string, deleted map[string]struct{}) error { + if len(deleted) == 0 { + return nil + } + scopes := []string{filepath.Join(memoryRoot, "global")} + if projectDir != "" { + scopes = append(scopes, filepath.Join(memoryRoot, "projects", projectID(projectDir))) + } + for _, scope := range scopes { + index := filepath.Join(scope, "MEMORY.md") + if _, err := os.Stat(index); err != nil { + if os.IsNotExist(err) { + continue + } + return err + } + if !withinScope(memoryRoot, projectDir, index) { + continue + } + tokens := indexRefTokens(memoryRoot, scope, deleted) + if len(tokens) == 0 { + continue + } + raw, err := os.ReadFile(index) + if err != nil { + return err + } + newBody, changed := stripDanglingIndexLines(string(raw), tokens) + if changed { + if err := atomicWrite(index, []byte(newBody)); err != nil { + return err + } + } + } + return nil +} + +// indexRefTokens is the set of substrings that identify a deleted file inside a +// MEMORY.md index line: its absolute path, its path relative to the memory root +// and to the scope root, and its bare basename. A line containing any of these +// is treated as a link/reference to the removed entry. MEMORY.md itself is never +// a token (it is never a consolidation deletion target). +func indexRefTokens(memoryRoot, scope string, deleted map[string]struct{}) map[string]struct{} { + tokens := make(map[string]struct{}) + for p := range deleted { + clean := filepath.Clean(p) + add := func(s string) { + if s != "" && s != "." { + tokens[filepath.ToSlash(s)] = struct{}{} + } + } + add(clean) + if rel, err := filepath.Rel(memoryRoot, clean); err == nil && !strings.HasPrefix(rel, "..") { + add(rel) + } + if rel, err := filepath.Rel(scope, clean); err == nil && !strings.HasPrefix(rel, "..") { + add(rel) + } + add(filepath.Base(clean)) + } + return tokens +} + +// stripDanglingIndexLines removes every line of body that references any of the +// reference tokens, returning the rewritten body and whether anything changed. +// It matches on the forward-slash form of each line so Windows-style separators +// in the index still match the slash tokens. Matching is boundary-aware: a token +// (e.g. the basename "b.md") only matches when it is not embedded inside a longer +// filename token (so "club.md" or "b.mdx" is not mistaken for "b.md"), avoiding +// dropping unrelated index lines. +func stripDanglingIndexLines(body string, tokens map[string]struct{}) (string, bool) { + lines := strings.Split(body, "\n") + kept := make([]string, 0, len(lines)) + changed := false + for _, line := range lines { + probe := filepath.ToSlash(line) + drop := false + for tok := range tokens { + if containsRefToken(probe, tok) { + drop = true + break + } + } + if drop { + changed = true + continue + } + kept = append(kept, line) + } + if !changed { + return body, false + } + return strings.Join(kept, "\n"), true +} + +// containsRefToken reports whether tok occurs in line at a filename boundary: +// the characters immediately before and after the match must not be filename +// continuation characters ([A-Za-z0-9_-]). This lets "b.md" match "user/b.md", +// "(b.md)" and "- b.md" while rejecting "club.md" and "b.mdx". +func containsRefToken(line, tok string) bool { + if tok == "" { + return false + } + from := 0 + for { + i := strings.Index(line[from:], tok) + if i < 0 { + return false + } + start := from + i + end := start + len(tok) + if !isFilenameChar(byteAt(line, start-1)) && !isFilenameChar(byteAt(line, end)) { + return true + } + from = start + 1 + } +} + +// byteAt returns line[i], or 0 when i is out of range (treated as a boundary). +func byteAt(line string, i int) byte { + if i < 0 || i >= len(line) { + return 0 + } + return line[i] +} + +// isFilenameChar reports whether b can appear inside a bare filename token, used +// to detect whether a reference-token match is embedded in a longer name. +func isFilenameChar(b byte) bool { + switch { + case b >= 'a' && b <= 'z', b >= 'A' && b <= 'Z', b >= '0' && b <= '9': + return true + case b == '_' || b == '-': + return true + } + return false +} + diff --git a/pigo/internal/dream/apply_test.go b/pigo/internal/dream/apply_test.go new file mode 100644 index 0000000..be8889b --- /dev/null +++ b/pigo/internal/dream/apply_test.go @@ -0,0 +1,143 @@ +package dream + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/memory" +) + +func TestUpdateScopeIndexesDropsDanglingLinks(t *testing.T) { + root := t.TempDir() + // A global MEMORY.md linking to two entries; b.md will be deleted. + writeMemFile(t, root, "global/user/a.md", "keep me") + b := writeMemFile(t, root, "global/user/b.md", "remove me") + idx := writeMemFile(t, root, "global/MEMORY.md", + "# Index\n- [a](user/a.md)\n- [b](user/b.md)\n- freeform note\n") + + deleted := map[string]struct{}{filepath.Clean(b): {}} + if err := updateScopeIndexes(root, "", deleted); err != nil { + t.Fatalf("updateScopeIndexes: %v", err) + } + raw, err := os.ReadFile(idx) + if err != nil { + t.Fatal(err) + } + got := string(raw) + if strings.Contains(got, "user/b.md") || strings.Contains(got, "b.md") { + t.Fatalf("dangling link to b.md not removed:\n%s", got) + } + if !strings.Contains(got, "user/a.md") { + t.Fatalf("live link to a.md wrongly removed:\n%s", got) + } + if !strings.Contains(got, "freeform note") { + t.Fatalf("unrelated line wrongly removed:\n%s", got) + } +} + +// TestUpdateScopeIndexesBoundaryMatch: deleting b.md must not drop an index line +// referencing a different entry whose name merely contains "b.md" as a +// substring (e.g. club.md). +func TestUpdateScopeIndexesBoundaryMatch(t *testing.T) { + root := t.TempDir() + writeMemFile(t, root, "global/user/club.md", "keep me") + b := writeMemFile(t, root, "global/user/b.md", "remove me") + idx := writeMemFile(t, root, "global/MEMORY.md", + "- [club](user/club.md)\n- [b](user/b.md)\n") + + deleted := map[string]struct{}{filepath.Clean(b): {}} + if err := updateScopeIndexes(root, "", deleted); err != nil { + t.Fatalf("updateScopeIndexes: %v", err) + } + got, _ := os.ReadFile(idx) + if !strings.Contains(string(got), "user/club.md") { + t.Fatalf("club.md link wrongly removed by substring match:\n%s", got) + } + if strings.Contains(string(got), "user/b.md") { + t.Fatalf("b.md link not removed:\n%s", got) + } +} + +func TestUpdateScopeIndexesNoIndexIsNoOp(t *testing.T) { + root := t.TempDir() + writeMemFile(t, root, "global/user/a.md", "x") + deleted := map[string]struct{}{filepath.Join(root, "global", "user", "a.md"): {}} + if err := updateScopeIndexes(root, "", deleted); err != nil { + t.Fatalf("updateScopeIndexes with no MEMORY.md should be a no-op, got %v", err) + } +} + +// TestRunEndToEndWithMerge exercises the full non-dry-run write path with a stub +// Consolidator that merges b.md into a.md and prunes c.md: files converge on +// disk, the MEMORY.md index loses its dangling links, Reconcile runs, the Report +// counters are correct, and a full-text search no longer hits the merged-away +// fragment (US-003 / US-006 / US-009). +func TestRunEndToEndWithMerge(t *testing.T) { + root := t.TempDir() + a := writeMemFile(t, root, "global/user/a.md", "shared topic original a") + b := writeMemFile(t, root, "global/user/b.md", "shared topic zebrafragment only in b") + c := writeMemFile(t, root, "global/user/c.md", "outdated standalone note") + idx := writeMemFile(t, root, "global/MEMORY.md", + "# Index\n- [a](user/a.md)\n- [b](user/b.md)\n- [c](user/c.md)\n") + + stub := &stubConsolidator{result: ConsolidateResult{ + MergedBodies: map[string]string{a: "shared topic merged and current"}, + Deletions: []string{b, c}, + Merged: 1, + Pruned: 1, + Notes: []string{"pruned c: outdated"}, + }} + r := &Runner{MemoryRoot: root, Consolidator: stub} + rep, err := r.Run(context.Background(), RunOptions{}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if !stub.called { + t.Fatal("Consolidator not called") + } + if rep.Merged != 1 || rep.Pruned != 1 { + t.Fatalf("counters wrong: %+v", rep) + } + // a.md rewritten, b.md + c.md gone. + if got, _ := os.ReadFile(a); string(got) != "shared topic merged and current" { + t.Fatalf("a.md not rewritten: %q", got) + } + if _, err := os.Stat(b); !os.IsNotExist(err) { + t.Fatal("b.md should be deleted") + } + if _, err := os.Stat(c); !os.IsNotExist(err) { + t.Fatal("c.md should be deleted") + } + // MEMORY.md no longer links to the removed entries. + rawIdx, _ := os.ReadFile(idx) + if strings.Contains(string(rawIdx), "b.md") || strings.Contains(string(rawIdx), "c.md") { + t.Fatalf("MEMORY.md retains dangling links:\n%s", rawIdx) + } + if !strings.Contains(string(rawIdx), "a.md") { + t.Fatalf("MEMORY.md lost the live a.md link:\n%s", rawIdx) + } + // Reconcile ran and indexed the surviving files. + if rep.Reconciled.Indexed == 0 { + t.Fatalf("Reconcile did not index: %+v", rep.Reconciled) + } + if rep.FilesAfter != 2 { // a.md + MEMORY.md + t.Fatalf("FilesAfter = %d, want 2", rep.FilesAfter) + } + + // The merged-away fragment must no longer be searchable. + store, err := memory.Open(filepath.Join(root, "index.db"), root, "") + if err != nil { + t.Fatalf("open store: %v", err) + } + defer store.Close() + hits, err := store.Search("zebrafragment", memory.SearchOptions{ReconcileFirst: true}) + if err != nil { + t.Fatalf("search: %v", err) + } + if len(hits) != 0 { + t.Fatalf("merged-away fragment still searchable: %+v", hits) + } +} diff --git a/pigo/internal/dream/config.go b/pigo/internal/dream/config.go new file mode 100644 index 0000000..02de592 --- /dev/null +++ b/pigo/internal/dream/config.go @@ -0,0 +1,48 @@ +// Package dream implements the /dream memory-consolidation feature's foundation +// layer: the resolved [dream] configuration, on-disk run state, and the +// deterministic due-check that decides whether an auto-trigger is warranted. +// This layer has no LLM dependency; the runner, scheduler, lock, and plan/apply +// logic live in later nodes. See tasks/spec-dream-memory-consolidation.md +// §3.2/§3.3/§5.1/§5.4. +package dream + +// Built-in defaults for the [dream] table, exposed so config normalization and +// tests share one source of truth. +const ( + DefaultEnabled = true + DefaultIntervalDays = 7 + DefaultRecentSessions = 20 +) + +// Config is the resolved [dream] configuration with defaults applied. It is the +// shape the scheduler and runner consume, distinct from the raw +// config.DreamConfig (which uses *bool / zero to distinguish "unset"). +type Config struct { + Enabled bool + IntervalDays int + RecentSessions int +} + +// NewConfig normalizes a raw [dream] table into a Config, applying defaults: a +// nil enabled pointer means true (only an explicit false disables); a +// non-positive interval_days falls back to 7; a non-positive recent_sessions +// falls back to 20. A missing [dream] table is representable as the zero +// arguments (nil, 0, 0) and yields all defaults, so parsing never errors on an +// absent table. +func NewConfig(enabled *bool, intervalDays, recentSessions int) Config { + c := Config{ + Enabled: DefaultEnabled, + IntervalDays: intervalDays, + RecentSessions: recentSessions, + } + if enabled != nil { + c.Enabled = *enabled + } + if c.IntervalDays <= 0 { + c.IntervalDays = DefaultIntervalDays + } + if c.RecentSessions <= 0 { + c.RecentSessions = DefaultRecentSessions + } + return c +} diff --git a/pigo/internal/dream/config_test.go b/pigo/internal/dream/config_test.go new file mode 100644 index 0000000..e8787ef --- /dev/null +++ b/pigo/internal/dream/config_test.go @@ -0,0 +1,61 @@ +package dream + +import "testing" + +func boolPtr(b bool) *bool { return &b } + +func TestNewConfigDefaults(t *testing.T) { + // Missing [dream] table: nil enabled, zero ints → all defaults. + c := NewConfig(nil, 0, 0) + if !c.Enabled { + t.Errorf("Enabled = false, want true (nil → true)") + } + if c.IntervalDays != DefaultIntervalDays { + t.Errorf("IntervalDays = %d, want %d", c.IntervalDays, DefaultIntervalDays) + } + if c.RecentSessions != DefaultRecentSessions { + t.Errorf("RecentSessions = %d, want %d", c.RecentSessions, DefaultRecentSessions) + } +} + +func TestNewConfigEnabledSemantics(t *testing.T) { + tests := []struct { + name string + enabled *bool + want bool + }{ + {"nil treated as true", nil, true}, + {"explicit true", boolPtr(true), true}, + {"explicit false", boolPtr(false), false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := NewConfig(tt.enabled, 0, 0).Enabled; got != tt.want { + t.Errorf("Enabled = %v, want %v", got, tt.want) + } + }) + } +} + +func TestNewConfigNonPositiveFallback(t *testing.T) { + tests := []struct { + name string + interval, recent int + wantInterval, wantRcnt int + }{ + {"zero falls back", 0, 0, DefaultIntervalDays, DefaultRecentSessions}, + {"negative falls back", -3, -1, DefaultIntervalDays, DefaultRecentSessions}, + {"positive preserved", 14, 50, 14, 50}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := NewConfig(nil, tt.interval, tt.recent) + if c.IntervalDays != tt.wantInterval { + t.Errorf("IntervalDays = %d, want %d", c.IntervalDays, tt.wantInterval) + } + if c.RecentSessions != tt.wantRcnt { + t.Errorf("RecentSessions = %d, want %d", c.RecentSessions, tt.wantRcnt) + } + }) + } +} diff --git a/pigo/internal/dream/consolidator.go b/pigo/internal/dream/consolidator.go new file mode 100644 index 0000000..af211e2 --- /dev/null +++ b/pigo/internal/dream/consolidator.go @@ -0,0 +1,434 @@ +package dream + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "sort" + "strings" +) + +// This file implements the real, LLM-backed Consolidator (SPEC §5.1 step 5, +// §5.1.1). It turns the deterministic Plan into a prompt for the main-session +// model, asks it to confirm semantic merges and conservative prunes, and parses +// the strict-JSON response back into a ConsolidateResult the Runner applies. +// +// The deterministic half (exact dedup, dead-path cleanup, MEMORY.md rewrite, +// Reconcile, counters) stays in the Runner; this file is purely the semantic +// LLM step. It never touches disk — it only produces decisions the Runner +// path-guards and applies. + +// completeFn performs a single LLM completion: given the system and user +// prompts it returns the model's text response, or an error for a hard +// transport/provider failure. The provider-backed implementation lives in +// apply.go (newModelCompleter); tests inject a canned function so no live model +// is ever called. +type completeFn func(ctx context.Context, systemPrompt, userPrompt string) (string, error) + +// defaultBodyBudget caps how many bytes of each entry body are shown to the +// model, bounding prompt size on large memory libraries (SPEC §8.2 token +// budget). Bodies longer than this are truncated with a marker so the model +// still sees the leading, usually most salient, content. +const defaultBodyBudget = 6000 + +// llmConsolidator is the production Consolidator: it drives the main-session +// model through complete and parses the response. bodyBudget (0 → default) +// bounds per-entry body size in the prompt. +type llmConsolidator struct { + complete completeFn + bodyBudget int +} + +// Consolidate builds the prompt from the plan, runs one model completion, and +// parses the response into merge/prune decisions. A hard model/transport error +// is returned (the Runner maps it to a failed run, SPEC §5.5). A well-formed +// call whose text cannot be parsed is NOT an error: it yields an empty result +// with an explanatory note, so an unparseable response conservatively KEEPs +// everything (PRD FR-14) and the deterministic pass still applies. +func (c *llmConsolidator) Consolidate(ctx context.Context, in ConsolidateInput) (ConsolidateResult, error) { + if c.complete == nil { + return ConsolidateResult{}, fmt.Errorf("dream: llmConsolidator has no completion function") + } + + var res ConsolidateResult + eligible := eligibleFiles(in.Plan) + if len(eligible) > 0 { + budget := c.bodyBudget + if budget <= 0 { + budget = defaultBodyBudget + } + prompt := buildConsolidatePrompt(in, eligible, budget) + + raw, err := c.complete(ctx, dreamSystemPrompt, prompt) + if err != nil { + return ConsolidateResult{}, fmt.Errorf("dream: model completion: %w", err) + } + + allowed := make(map[string]struct{}, len(eligible)) + for _, f := range eligible { + allowed[filepath.Clean(f.Path)] = struct{}{} + } + res = parseConsolidateResponse(raw, allowed) + } + + // Distillation pass (SPEC §5.3, PRD US-005/FR-13): a SEPARATE model call over + // the recent-session transcripts the Runner collected. It runs even when the + // library is empty (nothing to merge/prune) so a first-time distill can seed + // memory from sessions. A hard model failure aborts the run; a well-formed + // call yielding nothing simply adds no entries (Runner records the no-op). + if strings.TrimSpace(in.Transcripts) != "" { + if err := c.distill(ctx, in, &res); err != nil { + return ConsolidateResult{}, err + } + } + return res, nil +} + +// distill runs the JSONL distillation model call and folds the resulting new +// entries into res: it prompts the distiller with the transcripts plus a summary +// of the existing library (so the model avoids re-proposing known facts), parses +// the response into path-guarded NewEntry writes deduped against the existing +// memory, and bumps res.Distilled by the number added. A hard model/transport +// error is returned so the Runner marks the run failed (SPEC §5.5); an +// unparseable or empty response adds nothing and is not an error (conservative +// KEEP, PRD FR-14). +func (c *llmConsolidator) distill(ctx context.Context, in ConsolidateInput, res *ConsolidateResult) error { + prompt := buildDistillPrompt(in) + raw, err := c.complete(ctx, dreamDistillSystemPrompt, prompt) + if err != nil { + return fmt.Errorf("dream: distill completion: %w", err) + } + entries, notes := parseDistillResponse(raw, in.Plan.Files, in.MemoryRoot, in.ProjectDir) + res.NewEntries = append(res.NewEntries, entries...) + res.Distilled += len(entries) + res.Notes = append(res.Notes, notes...) + return nil +} + +// buildDistillPrompt renders the distiller's user prompt: the target scope, a +// compact list of the titles/paths of existing memories (so the model does not +// re-propose known facts), and the recent-session transcripts. Existing bodies +// are summarized (path + leading text) rather than dumped in full to keep the +// prompt bounded; the Go side still enforces near-duplicate rejection. +func buildDistillPrompt(in ConsolidateInput) string { + var b strings.Builder + b.WriteString("# Memory distillation request\n\n") + if in.ProjectDir != "" { + b.WriteString("Current project: ") + b.WriteString(in.ProjectDir) + b.WriteByte('\n') + } else { + b.WriteString("Scope: global only\n") + } + + existing := eligibleFiles(in.Plan) + if len(existing) > 0 { + b.WriteString(fmt.Sprintf("\n## Existing memories (%d) — do NOT re-propose these\n\n", len(existing))) + for _, f := range existing { + b.WriteString("- ") + b.WriteString(clip(strings.TrimSpace(firstNonEmptyLine(f.Body)), 120)) + b.WriteByte('\n') + } + } + + b.WriteString("\n## Recent session transcripts\n\n") + b.WriteString(in.Transcripts) + b.WriteString("\n\nReturn the JSON object described in your instructions. Extract only genuinely new, durable facts. When in doubt, return no entries.\n") + return b.String() +} + +// firstNonEmptyLine returns the first non-blank line of s (trimmed), or "" when +// s is entirely blank. Used to label an existing memory in the distill prompt. +func firstNonEmptyLine(s string) string { + for _, line := range strings.Split(s, "\n") { + if t := strings.TrimSpace(line); t != "" { + return t + } + } + return "" +} + +// eligibleFiles is the subset of plan files the model may act on: it excludes +// MEMORY.md index files (they are indexes, not entries — #521 NEW_WORK: we +// special-case MEMORY.md out of merge/prune so the index is never folded into an +// entry) and the redundant members of an exact-dedupe group (g.Paths[1:], which +// the deterministic pass removes anyway — offering them would let the model +// merge into a path about to be deleted). The representative g.Paths[0] stays. +func eligibleFiles(plan Plan) []MemoryFile { + drop := make(map[string]struct{}) + for _, g := range plan.DedupeGroups { + for _, p := range g.Paths[1:] { + drop[filepath.Clean(p)] = struct{}{} + } + } + var out []MemoryFile + for _, f := range plan.Files { + if isMemoryIndex(f.Path) { + continue + } + if _, dup := drop[filepath.Clean(f.Path)]; dup { + continue + } + out = append(out, f) + } + return out +} + +// isMemoryIndex reports whether path is a scope MEMORY.md index file, which the +// consolidation step must never merge, rewrite, or prune. +func isMemoryIndex(path string) bool { + return strings.EqualFold(filepath.Base(path), "MEMORY.md") +} + +// buildConsolidatePrompt renders the user prompt: the scope, the eligible +// entries (path + scope/type + body, truncated to budget), and the deterministic +// hints (near-dup candidate pairs, dead local-path references). It only lists +// paths that are eligible, so the model is naturally steered away from MEMORY.md +// and duplicate paths. +func buildConsolidatePrompt(in ConsolidateInput, eligible []MemoryFile, budget int) string { + var b strings.Builder + b.WriteString("# Memory consolidation request\n\n") + b.WriteString("Memory root: ") + b.WriteString(in.MemoryRoot) + b.WriteByte('\n') + if in.ProjectDir != "" { + b.WriteString("Active project scope: ") + b.WriteString(in.ProjectDir) + b.WriteByte('\n') + } else { + b.WriteString("Scope: global only\n") + } + b.WriteString(fmt.Sprintf("\n## Entries (%d)\n\n", len(eligible))) + for i, f := range eligible { + typ := f.Type + if typ == "" { + typ = "(root)" + } + b.WriteString(fmt.Sprintf("### [%d] %s\n", i+1, f.Path)) + b.WriteString(fmt.Sprintf("scope=%s type=%s bytes=%d\n\n", f.Scope, typ, f.Size)) + b.WriteString("```\n") + b.WriteString(truncateBody(f.Body, budget)) + b.WriteString("\n```\n\n") + } + + if pairs := eligiblePairs(in.Plan, eligible); len(pairs) > 0 { + b.WriteString("## Near-duplicate candidate pairs (merge only if truly overlapping)\n\n") + for _, p := range pairs { + b.WriteString(fmt.Sprintf("- %s <-> %s (similarity %.2f)\n", p.A, p.B, p.Similarity)) + } + b.WriteByte('\n') + } + + if refs := eligibleInvalidRefs(in.Plan, eligible); len(refs) > 0 { + b.WriteString("## Entries referencing local files that no longer exist\n") + b.WriteString("(the dead reference text is cleaned automatically; only PRUNE an entry if losing that reference leaves it meaningless)\n\n") + for _, r := range refs { + b.WriteString(fmt.Sprintf("- %s references missing %s\n", r.File, r.Ref)) + } + b.WriteByte('\n') + } + + b.WriteString("Return the JSON object described in your instructions. When in doubt, KEEP.\n") + return b.String() +} + +// truncateBody trims body to at most budget bytes on a rune boundary, appending +// a marker when truncation occurred so the model knows content was elided. +func truncateBody(body string, budget int) string { + if budget <= 0 || len(body) <= budget { + return body + } + cut := budget + for cut > 0 && !isRuneStart(body[cut]) { + cut-- + } + return body[:cut] + "\n…[truncated]" +} + +// isRuneStart reports whether b is not a UTF-8 continuation byte, so a truncation +// cut there does not split a multi-byte rune. +func isRuneStart(b byte) bool { return b&0xC0 != 0x80 } + +// eligiblePairs filters the plan's near-dup pairs to those whose BOTH members +// are eligible (both are still offered to the model), so we never point the +// model at a MEMORY.md or a to-be-deduped duplicate. +func eligiblePairs(plan Plan, eligible []MemoryFile) []NearDupPair { + ok := make(map[string]struct{}, len(eligible)) + for _, f := range eligible { + ok[filepath.Clean(f.Path)] = struct{}{} + } + var out []NearDupPair + for _, p := range plan.NearDupPairs { + _, a := ok[filepath.Clean(p.A)] + _, b := ok[filepath.Clean(p.B)] + if a && b { + out = append(out, p) + } + } + return out +} + +// eligibleInvalidRefs filters dead-path references to those on eligible files. +func eligibleInvalidRefs(plan Plan, eligible []MemoryFile) []InvalidPathRef { + ok := make(map[string]struct{}, len(eligible)) + for _, f := range eligible { + ok[filepath.Clean(f.Path)] = struct{}{} + } + var out []InvalidPathRef + for _, r := range plan.InvalidPathRefs { + if _, found := ok[filepath.Clean(r.File)]; found { + out = append(out, r) + } + } + return out +} + +// modelDecision mirrors the strict JSON output schema the dream prompt requires. +type modelDecision struct { + Merges []struct { + Keep string `json:"keep"` + Body string `json:"body"` + Remove []string `json:"remove"` + } `json:"merges"` + Prunes []struct { + Path string `json:"path"` + Reason string `json:"reason"` + } `json:"prunes"` + Notes []string `json:"notes"` +} + +// parseConsolidateResponse turns the model's text into a ConsolidateResult, +// validating every path against allowed (the eligible input paths). It is +// conservative by construction: any decision that references an unknown path, a +// MEMORY.md index, an empty merged body, or an empty prune reason is dropped +// rather than obeyed, and an unparseable response yields an empty result with a +// note (PRD FR-14 — default to KEEP on uncertainty). It never returns an error; +// callers treat a hard model failure separately. +func parseConsolidateResponse(raw string, allowed map[string]struct{}) ConsolidateResult { + body := extractJSONObject(raw) + if body == "" { + return ConsolidateResult{Notes: []string{"dream: model response contained no JSON object; kept all entries"}} + } + var dec modelDecision + if err := json.Unmarshal([]byte(body), &dec); err != nil { + return ConsolidateResult{Notes: []string{"dream: model response was not valid JSON; kept all entries"}} + } + + var res ConsolidateResult + res.MergedBodies = make(map[string]string) + delSet := make(map[string]struct{}) + + valid := func(p string) (string, bool) { + clean := filepath.Clean(strings.TrimSpace(p)) + if clean == "" || clean == "." { + return "", false + } + if isMemoryIndex(clean) { + return "", false + } + if _, ok := allowed[clean]; !ok { + return "", false + } + return clean, true + } + + for _, m := range dec.Merges { + keep, ok := valid(m.Keep) + if !ok || strings.TrimSpace(m.Body) == "" { + // Unknown/invalid target or an empty rewrite: skip, keep everything. + continue + } + removed := 0 + for _, r := range m.Remove { + rp, ok := valid(r) + if !ok || rp == keep { + continue + } + if _, dup := delSet[rp]; dup { + continue + } + delSet[rp] = struct{}{} + removed++ + } + if removed == 0 { + // A merge that removes nothing is a no-op rewrite; ignore it to avoid + // gratuitously touching a file the model merely echoed back. + continue + } + res.MergedBodies[keep] = m.Body + res.Merged += removed + } + + for _, p := range dec.Prunes { + pp, ok := valid(p.Path) + if !ok || strings.TrimSpace(p.Reason) == "" { + // No path or no stated reason → conservative KEEP. + continue + } + if _, dup := delSet[pp]; dup { + continue + } + // Never prune an entry we are simultaneously keeping as a merge target. + if _, kept := res.MergedBodies[pp]; kept { + continue + } + delSet[pp] = struct{}{} + res.Pruned++ + res.Notes = append(res.Notes, fmt.Sprintf("pruned %s: %s", pp, strings.TrimSpace(p.Reason))) + } + + if len(res.MergedBodies) == 0 { + res.MergedBodies = nil + } + for p := range delSet { + res.Deletions = append(res.Deletions, p) + } + sort.Strings(res.Deletions) + + for _, n := range dec.Notes { + if s := strings.TrimSpace(n); s != "" { + res.Notes = append(res.Notes, s) + } + } + return res +} + +// extractJSONObject returns the outermost {...} span of s, tolerating models +// that wrap the object in prose or Markdown code fences. It returns "" when no +// balanced object is found. +func extractJSONObject(s string) string { + start := strings.IndexByte(s, '{') + if start < 0 { + return "" + } + depth := 0 + inStr := false + esc := false + for i := start; i < len(s); i++ { + ch := s[i] + if inStr { + switch { + case esc: + esc = false + case ch == '\\': + esc = true + case ch == '"': + inStr = false + } + continue + } + switch ch { + case '"': + inStr = true + case '{': + depth++ + case '}': + depth-- + if depth == 0 { + return s[start : i+1] + } + } + } + return "" +} diff --git a/pigo/internal/dream/consolidator_test.go b/pigo/internal/dream/consolidator_test.go new file mode 100644 index 0000000..2804af3 --- /dev/null +++ b/pigo/internal/dream/consolidator_test.go @@ -0,0 +1,195 @@ +package dream + +import ( + "context" + "errors" + "path/filepath" + "strings" + "testing" +) + +// canned builds an allowed-set + Plan pair for parser tests from a list of +// absolute paths. +func allowedSet(paths ...string) map[string]struct{} { + m := make(map[string]struct{}, len(paths)) + for _, p := range paths { + m[filepath.Clean(p)] = struct{}{} + } + return m +} + +func TestBuildConsolidatePromptListsEntriesAndHints(t *testing.T) { + root := "/mem" + a := "/mem/global/user/a.md" + b := "/mem/global/user/b.md" + idx := "/mem/global/MEMORY.md" + plan := Plan{ + Files: []MemoryFile{ + {Path: a, Scope: "global", Type: "user", Size: 5, Body: "alpha body"}, + {Path: b, Scope: "global", Type: "user", Size: 5, Body: "beta body"}, + {Path: idx, Scope: "global", Type: "", Size: 3, Body: "- [a](user/a.md)"}, + }, + NearDupPairs: []NearDupPair{{A: a, B: b, Similarity: 0.82}}, + InvalidPathRefs: []InvalidPathRef{{File: a, Ref: "./gone.go"}}, + } + eligible := eligibleFiles(plan) + if len(eligible) != 2 { + t.Fatalf("eligibleFiles = %d, want 2 (MEMORY.md excluded)", len(eligible)) + } + prompt := buildConsolidatePrompt(ConsolidateInput{Plan: plan, MemoryRoot: root, ProjectDir: ""}, eligible, defaultBodyBudget) + + for _, want := range []string{a, b, "alpha body", "beta body", "similarity 0.82", "./gone.go", "global only"} { + if !strings.Contains(prompt, want) { + t.Errorf("prompt missing %q\n---\n%s", want, prompt) + } + } + if strings.Contains(prompt, "MEMORY.md") { + t.Errorf("prompt must not offer the MEMORY.md index as an entry:\n%s", prompt) + } +} + +func TestTruncateBody(t *testing.T) { + if got := truncateBody("short", 100); got != "short" { + t.Fatalf("no truncation expected, got %q", got) + } + long := strings.Repeat("x", 50) + got := truncateBody(long, 10) + if !strings.HasPrefix(got, strings.Repeat("x", 10)) || !strings.Contains(got, "truncated") { + t.Fatalf("truncateBody = %q", got) + } +} + +func TestParseConsolidateResponseMergeAndPrune(t *testing.T) { + a := "/mem/global/user/a.md" + b := "/mem/global/user/b.md" + c := "/mem/global/user/c.md" + allowed := allowedSet(a, b, c) + + raw := "here you go:\n```json\n" + `{ + "merges": [{"keep": "` + a + `", "body": "merged body", "remove": ["` + b + `"]}], + "prunes": [{"path": "` + c + `", "reason": "superseded by newer note"}], + "notes": ["did the thing"] + }` + "\n```\n" + + res := parseConsolidateResponse(raw, allowed) + if res.Merged != 1 { + t.Errorf("Merged = %d, want 1", res.Merged) + } + if res.Pruned != 1 { + t.Errorf("Pruned = %d, want 1", res.Pruned) + } + if got := res.MergedBodies[a]; got != "merged body" { + t.Errorf("MergedBodies[a] = %q", got) + } + wantDel := map[string]bool{b: true, c: true} + if len(res.Deletions) != 2 { + t.Fatalf("Deletions = %v, want b and c", res.Deletions) + } + for _, d := range res.Deletions { + if !wantDel[d] { + t.Errorf("unexpected deletion %q", d) + } + } + joined := strings.Join(res.Notes, "|") + if !strings.Contains(joined, "superseded by newer note") || !strings.Contains(joined, "did the thing") { + t.Errorf("notes missing content: %v", res.Notes) + } +} + +func TestParseConsolidateResponseRejectsUnknownAndUnsafePaths(t *testing.T) { + a := "/mem/global/user/a.md" + allowed := allowedSet(a) + raw := `{ + "merges": [{"keep": "/etc/passwd", "body": "x", "remove": ["` + a + `"]}], + "prunes": [{"path": "/mem/global/MEMORY.md", "reason": "index"}] + }` + res := parseConsolidateResponse(raw, allowed) + if res.Merged != 0 || res.Pruned != 0 || len(res.Deletions) != 0 { + t.Fatalf("unsafe/unknown paths must be ignored, got %+v", res) + } +} + +func TestParseConsolidateResponseConservativeOnEmptyReasonAndBadJSON(t *testing.T) { + a := "/mem/global/user/a.md" + allowed := allowedSet(a) + + // Empty prune reason → KEEP. + res := parseConsolidateResponse(`{"prunes":[{"path":"`+a+`","reason":""}]}`, allowed) + if res.Pruned != 0 || len(res.Deletions) != 0 { + t.Fatalf("empty reason must KEEP, got %+v", res) + } + + // Unparseable → empty result with a note, never a deletion. + res = parseConsolidateResponse("the model rambled with no json", allowed) + if res.Merged != 0 || res.Pruned != 0 || len(res.Deletions) != 0 { + t.Fatalf("bad JSON must KEEP everything, got %+v", res) + } + if len(res.Notes) == 0 { + t.Fatal("expected an explanatory note on unparseable response") + } +} + +func TestParseConsolidateResponseNoOpMergeIgnored(t *testing.T) { + a := "/mem/global/user/a.md" + allowed := allowedSet(a) + // A merge that removes nothing must not touch the file. + res := parseConsolidateResponse(`{"merges":[{"keep":"`+a+`","body":"rewritten","remove":[]}]}`, allowed) + if len(res.MergedBodies) != 0 || res.Merged != 0 { + t.Fatalf("no-op merge must be ignored, got %+v", res) + } +} + +func TestLLMConsolidatorUsesCompleter(t *testing.T) { + root := "/mem" + a := "/mem/global/user/a.md" + b := "/mem/global/user/b.md" + plan := Plan{Files: []MemoryFile{ + {Path: a, Scope: "global", Type: "user", Body: "one"}, + {Path: b, Scope: "global", Type: "user", Body: "two"}, + }} + + var gotSystem, gotUser string + c := &llmConsolidator{complete: func(_ context.Context, sys, user string) (string, error) { + gotSystem, gotUser = sys, user + return `{"merges":[{"keep":"` + a + `","body":"merged","remove":["` + b + `"]}]}`, nil + }} + res, err := c.Consolidate(context.Background(), ConsolidateInput{Plan: plan, MemoryRoot: root}) + if err != nil { + t.Fatalf("Consolidate: %v", err) + } + if gotSystem != dreamSystemPrompt { + t.Error("system prompt not passed through") + } + if !strings.Contains(gotUser, a) { + t.Error("user prompt missing entry path") + } + if res.Merged != 1 || res.MergedBodies[a] != "merged" { + t.Fatalf("merge not applied: %+v", res) + } +} + +func TestLLMConsolidatorPropagatesHardError(t *testing.T) { + plan := Plan{Files: []MemoryFile{{Path: "/mem/global/user/a.md", Scope: "global", Type: "user", Body: "x"}}} + c := &llmConsolidator{complete: func(context.Context, string, string) (string, error) { + return "", errors.New("upstream 500") + }} + if _, err := c.Consolidate(context.Background(), ConsolidateInput{Plan: plan, MemoryRoot: "/mem"}); err == nil { + t.Fatal("expected hard completion error to propagate") + } +} + +func TestLLMConsolidatorSkipsWhenNoEligibleFiles(t *testing.T) { + called := false + c := &llmConsolidator{complete: func(context.Context, string, string) (string, error) { + called = true + return "{}", nil + }} + // Only a MEMORY.md index → nothing eligible → no model call. + plan := Plan{Files: []MemoryFile{{Path: "/mem/global/MEMORY.md", Scope: "global", Body: "idx"}}} + if _, err := c.Consolidate(context.Background(), ConsolidateInput{Plan: plan, MemoryRoot: "/mem"}); err != nil { + t.Fatalf("Consolidate: %v", err) + } + if called { + t.Fatal("model should not be called when no eligible files exist") + } +} diff --git a/pigo/internal/dream/distill.go b/pigo/internal/dream/distill.go new file mode 100644 index 0000000..e9039b3 --- /dev/null +++ b/pigo/internal/dream/distill.go @@ -0,0 +1,411 @@ +package dream + +// This file implements the JSONL distillation step of a /dream run (SPEC §5.3, +// PRD US-005 / FR-13): it selects the current project's recent sessions, reads +// their transcripts (truncated to a token/byte budget, SPEC §8.2), and — via the +// LLM distiller in consolidator.go — turns durable facts into new memory +// entries, deduped against the existing library so nothing already recorded is +// re-added. +// +// The deterministic half lives here (session-window selection, project filter, +// transcript rendering, dedup, path construction); the semantic half (which +// facts are durable) is a separate LLM call driven by the llmConsolidator. This +// keeps the merge/prune pass and its parser (consolidator.go) untouched: the +// Runner gathers the transcripts, the consolidator runs a distinct distill +// completion with its own prompt/schema, and the distilled NewEntries are folded +// into the same ConsolidateResult the Runner already applies (wiring choice "b" +// from #523: a separate step appended to the result, chosen over folding it into +// the merge/prune prompt so each concern keeps its own input, prompt and schema +// and the well-tested merge/prune parser is not perturbed). + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/session" +) + +// defaultTranscriptBudget bounds the total bytes of session transcript text fed +// to the distiller, keeping the distill prompt within the model context even on +// long or numerous recent sessions (SPEC §8.2 token budget). Sessions are added +// most-recent-first until the budget is reached. +const defaultTranscriptBudget = 24000 + +// distillNewEntryTypes is the closed set of durable memory types the distiller +// may emit (PRD FR-13: user / feedback / project / reference). Ephemeral kinds +// (checkpoint / progress) are intentionally excluded — one-shot task state is +// not distilled into long-term memory. +var distillNewEntryTypes = map[string]struct{}{ + "user": {}, + "feedback": {}, + "project": {}, + "reference": {}, +} + +// SessionSource is the read-only view of the session store the distiller needs: +// list session headers and load a session's messages. *session.Store satisfies +// it; tests inject a stub so no real session files (or LLM) are required. +type SessionSource interface { + List() ([]session.SessionHeader, error) + Load(id string) (session.SessionHeader, agentcore.MessageList, error) +} + +// resolveSessionStore opens the default session store rooted at +// $PIGO_HOME/sessions (else ~/.pigo/sessions), mirroring +// headless.SessionStore's resolution. It is duplicated here rather than imported +// to keep internal/dream free of a dependency on the CLI assembly layer, exactly +// as ResolveMemoryRoot duplicates the memory-root resolution. A resolution +// failure yields a nil source and no error: distillation then degrades to a +// no-op rather than failing the whole dream run. +func resolveSessionStore() (SessionSource, error) { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return nil, nil + } + dir = filepath.Join(home, ".pigo") + } + store, err := session.NewStore(filepath.Join(dir, "sessions")) + if err != nil { + return nil, nil + } + return store, nil +} + +// sessionMatchesProject reports whether a session belongs to the project rooted +// at projectDir. Attribution is by the session's recorded Cwd: two directories +// belong to the same project when their stable project ids match (the same id +// internal/memory and BuildPlan derive). A session with an empty Cwd is +// unattributed and never matches; an empty projectDir (a global-only run) has no +// project to match, so nothing is selected. +func sessionMatchesProject(h session.SessionHeader, projectDir string) bool { + if projectDir == "" || h.Cwd == "" { + return false + } + return projectID(h.Cwd) == projectID(projectDir) +} + +// collectRecentSessions selects the recent sessions to distill, applying the +// SPEC §5.3 combined window over the project-filtered set: +// - state.LastRunAt non-zero → every matching session updated strictly after +// the last run (incremental distillation since the last dream). +// - state.LastRunAt zero (never run) → the most-recent recentN matching +// sessions by UpdatedAt descending (a bounded first-run window). +// +// recentN falls back to DefaultRecentSessions when non-positive. The result is +// ordered most-recent-first so transcript budgeting keeps the freshest context. +func collectRecentSessions(sessions []session.SessionHeader, state State, projectDir string, recentN int) []session.SessionHeader { + if recentN <= 0 { + recentN = DefaultRecentSessions + } + var matched []session.SessionHeader + for _, h := range sessions { + if sessionMatchesProject(h, projectDir) { + matched = append(matched, h) + } + } + // Most-recent-first regardless of the source ordering. + sort.SliceStable(matched, func(i, j int) bool { + return matched[i].UpdatedAt.After(matched[j].UpdatedAt) + }) + + if !state.LastRunAt.IsZero() { + var out []session.SessionHeader + for _, h := range matched { + if h.UpdatedAt.After(state.LastRunAt) { + out = append(out, h) + } + } + return out + } + if len(matched) > recentN { + matched = matched[:recentN] + } + return matched +} + +// collectTranscripts selects the recent project sessions (collectRecentSessions) +// and renders their messages into a single transcript string truncated to +// budget bytes (SPEC §8.2). It returns "" when the source is nil or no session +// matches — the no-matching-sessions no-op (SPEC §5.5): the caller then records +// Distilled=0 with a "无新增" note. A per-session load error is skipped rather +// than failing the run (a single corrupt session must not abort distillation). +func collectTranscripts(src SessionSource, state State, projectDir string, recentN, budget int) string { + if src == nil { + return "" + } + if budget <= 0 { + budget = defaultTranscriptBudget + } + headers, err := src.List() + if err != nil { + return "" + } + selected := collectRecentSessions(headers, state, projectDir, recentN) + if len(selected) == 0 { + return "" + } + + var b strings.Builder + for _, h := range selected { + if b.Len() >= budget { + break + } + _, msgs, err := src.Load(h.ID) + if err != nil { + continue + } + section := renderTranscript(h, msgs) + if section == "" { + continue + } + remaining := budget - b.Len() + if len(section) > remaining { + section = truncateBody(section, remaining) + } + b.WriteString(section) + b.WriteString("\n") + } + return strings.TrimSpace(b.String()) +} + +// renderTranscript renders one session's messages as a compact, role-tagged +// transcript for the distiller. Tool-result bodies are included but bounded so a +// single huge tool output cannot dominate the budget; empty messages are +// skipped. The header line carries the session id and update time for context. +func renderTranscript(h session.SessionHeader, msgs agentcore.MessageList) string { + var b strings.Builder + b.WriteString(fmt.Sprintf("## Session %s (updated %s)\n", h.ID, h.UpdatedAt.UTC().Format("2006-01-02"))) + wrote := false + for _, m := range msgs { + var role, text string + switch mm := m.(type) { + case agentcore.UserMessage: + role, text = "user", agentcore.ContentToText(mm.Content) + case agentcore.AssistantMessage: + role, text = "assistant", agentcore.ContentToText(mm.Content) + case agentcore.ToolResultMessage: + role, text = "tool", clip(agentcore.ContentToText(mm.Content), 500) + case agentcore.CompactionMessage: + role, text = "summary", mm.Summary + default: + continue + } + text = strings.TrimSpace(text) + if text == "" { + continue + } + b.WriteString(role) + b.WriteString(": ") + b.WriteString(text) + b.WriteString("\n") + wrote = true + } + if !wrote { + return "" + } + return b.String() +} + +// clip trims s to at most n bytes on a rune boundary, appending an ellipsis when +// truncation occurred. Used to bound individual tool-result bodies inside a +// transcript so one large output does not crowd out the rest. +func clip(s string, n int) string { + if n <= 0 || len(s) <= n { + return s + } + cut := n + for cut > 0 && !isRuneStart(s[cut]) { + cut-- + } + return s[:cut] + "…" +} + +// distilledEntry mirrors one element of the distiller's strict-JSON output: a +// proposed new memory entry with its semantic type, target scope, a short title +// (turned into a filename slug) and the Markdown body to persist. +type distilledEntry struct { + Type string `json:"type"` + Scope string `json:"scope"` + Title string `json:"title"` + Body string `json:"body"` +} + +// distillResponse is the full strict-JSON schema the distill prompt requires. +type distillResponse struct { + Entries []distilledEntry `json:"entries"` + Notes []string `json:"notes"` +} + +// parseDistillResponse turns the distiller's text into concrete NewEntry writes, +// deduped against the existing memory bodies. It is conservative by +// construction: an unparseable response, an unknown/ephemeral type, an empty +// body, or an entry that is a near-duplicate of an existing memory (or of an +// already-accepted new entry) is dropped rather than written (PRD FR-13/FR-14). +// It never errors; a hard model failure is handled by the caller. +// +// memoryRoot + projectDir place each entry within the scope the Runner's +// withinScope guard permits: global entries under /global//, project +// entries under /projects///. A "project" entry on a global-only +// run (empty projectDir) has nowhere valid to live and is skipped. +func parseDistillResponse(raw string, existing []MemoryFile, memoryRoot, projectDir string) ([]NewEntry, []string) { + body := extractJSONObject(raw) + if body == "" { + return nil, nil + } + var resp distillResponse + if err := json.Unmarshal([]byte(body), &resp); err != nil { + return nil, nil + } + + // Precompute token sets of existing bodies for near-duplicate rejection. + existingTokens := make([]map[string]struct{}, 0, len(existing)) + for _, f := range existing { + existingTokens = append(existingTokens, tokenize(f.Body)) + } + + var out []NewEntry + var notes []string + acceptedTokens := make([]map[string]struct{}, 0) + usedPaths := make(map[string]struct{}) + + for _, e := range resp.Entries { + typ := strings.ToLower(strings.TrimSpace(e.Type)) + if _, ok := distillNewEntryTypes[typ]; !ok { + continue + } + body := strings.TrimSpace(e.Body) + if body == "" { + continue + } + scopeDir, ok := distillScopeDir(e.Scope, projectDir) + if !ok { + continue + } + tok := tokenize(body) + if isNearDup(tok, existingTokens) || isNearDup(tok, acceptedTokens) { + notes = append(notes, fmt.Sprintf("distill: skipped near-duplicate of existing memory (%s)", firstLine(e.Title, body))) + continue + } + path := distillPath(memoryRoot, scopeDir, typ, e.Title, body, usedPaths) + usedPaths[filepath.Clean(path)] = struct{}{} + acceptedTokens = append(acceptedTokens, tok) + out = append(out, NewEntry{Path: path, Body: ensureTrailingNewline(body)}) + } + + for _, n := range resp.Notes { + if s := strings.TrimSpace(n); s != "" { + notes = append(notes, s) + } + } + return out, notes +} + +// distillScopeDir maps a requested scope ("global"/"project") to its on-disk +// scope directory relative to the memory root. A "project" scope requires a +// project dir; without one the entry cannot be placed and ok is false. An empty +// or unrecognized scope defaults to global (the safe, always-valid target). +func distillScopeDir(scope, projectDir string) (string, bool) { + switch strings.ToLower(strings.TrimSpace(scope)) { + case "project", "projects": + if projectDir == "" { + return "", false + } + return filepath.Join("projects", projectID(projectDir)), true + default: + return "global", true + } +} + +// distillPath builds the absolute file path for a new distilled entry: +// ///.md. The slug derives from the title +// (sanitized) or, when absent, a short content hash, so a re-run of the same +// durable fact lands on a stable path (and is caught by the near-dup check +// against the now-existing file, keeping distillation idempotent). If the path +// is already taken within this run, a short body hash disambiguates it. +func distillPath(memoryRoot, scopeDir, typ, title, body string, used map[string]struct{}) string { + slug := slugify(title) + if slug == "" { + slug = shortHash(body) + } + name := slug + ".md" + path := filepath.Join(memoryRoot, scopeDir, typ, name) + if _, taken := used[filepath.Clean(path)]; taken { + path = filepath.Join(memoryRoot, scopeDir, typ, slug+"-"+shortHash(body)+".md") + } + return path +} + +// isNearDup reports whether tok is at or above NearDupThreshold Jaccard +// similarity with any set in others — the same conservative near-duplicate +// signal the deterministic plan uses for merge candidates, reused here so a +// distilled fact that already lives in memory is not re-added (PRD FR-13). +func isNearDup(tok map[string]struct{}, others []map[string]struct{}) bool { + for _, o := range others { + if jaccard(tok, o) >= NearDupThreshold { + return true + } + } + return false +} + +// slugify turns a title into a lowercase, dash-separated filename stem of ASCII +// alphanumerics, bounding the length so paths stay reasonable. Non-alphanumeric +// runs collapse to a single dash; leading/trailing dashes are trimmed. A title +// with no usable characters yields "". +func slugify(title string) string { + var b strings.Builder + lastDash := false + for _, r := range strings.ToLower(strings.TrimSpace(title)) { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9': + b.WriteRune(r) + lastDash = false + default: + if !lastDash && b.Len() > 0 { + b.WriteByte('-') + lastDash = true + } + } + } + slug := strings.Trim(b.String(), "-") + const maxLen = 60 + if len(slug) > maxLen { + slug = strings.Trim(slug[:maxLen], "-") + } + return slug +} + +// firstLine returns a short human-readable label for a note: the trimmed title +// when present, else the first line of body, truncated. +func firstLine(title, body string) string { + s := strings.TrimSpace(title) + if s == "" { + s = strings.TrimSpace(body) + } + if i := strings.IndexByte(s, '\n'); i >= 0 { + s = s[:i] + } + return clip(s, 60) +} + +// ensureTrailingNewline guarantees body ends with exactly one newline so written +// memory files match the one-entry-per-file convention and diff cleanly. +func ensureTrailingNewline(body string) string { + return strings.TrimRight(body, "\n") + "\n" +} + +// shortHash returns the first 8 hex chars of sha256(s), a stable disambiguator +// for slugs derived from identical/absent titles. +func shortHash(s string) string { + sum := sha256.Sum256([]byte(s)) + return hex.EncodeToString(sum[:])[:8] +} diff --git a/pigo/internal/dream/distill_test.go b/pigo/internal/dream/distill_test.go new file mode 100644 index 0000000..9b050bb --- /dev/null +++ b/pigo/internal/dream/distill_test.go @@ -0,0 +1,290 @@ +package dream + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/session" +) + +// stubSessions is an in-memory SessionSource for distill tests: it lists the +// headers it was given and returns canned messages per id, so no real session +// files (or LLM) are touched. +type stubSessions struct { + headers []session.SessionHeader + msgs map[string]agentcore.MessageList +} + +func (s *stubSessions) List() ([]session.SessionHeader, error) { return s.headers, nil } + +func (s *stubSessions) Load(id string) (session.SessionHeader, agentcore.MessageList, error) { + for _, h := range s.headers { + if h.ID == id { + return h, s.msgs[id], nil + } + } + return session.SessionHeader{}, nil, os.ErrNotExist +} + +func userMsg(text string) agentcore.UserMessage { + return agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent(text)}} +} + +func asstMsg(text string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent(text)}} +} + +// TestCollectRecentSessionsFirstRunWindow: never-run (zero LastRunAt) selects the +// most-recent recentN matching sessions, ordered most-recent-first, filtered to +// the active project. +func TestCollectRecentSessionsFirstRunWindow(t *testing.T) { + proj := t.TempDir() + other := t.TempDir() + base := time.Now().UTC() + headers := []session.SessionHeader{ + {ID: "s1", Cwd: proj, UpdatedAt: base.Add(-3 * time.Hour)}, + {ID: "s2", Cwd: proj, UpdatedAt: base.Add(-1 * time.Hour)}, + {ID: "s3", Cwd: other, UpdatedAt: base}, // different project, excluded + {ID: "s4", Cwd: proj, UpdatedAt: base.Add(-2 * time.Hour)}, + {ID: "s5", Cwd: "", UpdatedAt: base}, // unattributed, excluded + } + got := collectRecentSessions(headers, State{}, proj, 2) + if len(got) != 2 { + t.Fatalf("got %d sessions, want 2 (recentN cap)", len(got)) + } + if got[0].ID != "s2" || got[1].ID != "s4" { + t.Fatalf("wrong window/order: %s, %s (want s2, s4 most-recent-first)", got[0].ID, got[1].ID) + } +} + +// TestCollectRecentSessionsIncremental: a non-zero LastRunAt selects only +// matching sessions updated strictly after it (incremental distillation). +func TestCollectRecentSessionsIncremental(t *testing.T) { + proj := t.TempDir() + base := time.Now().UTC() + last := base.Add(-2 * time.Hour) + headers := []session.SessionHeader{ + {ID: "old", Cwd: proj, UpdatedAt: base.Add(-3 * time.Hour)}, // before last run + {ID: "new1", Cwd: proj, UpdatedAt: base.Add(-1 * time.Hour)}, + {ID: "new2", Cwd: proj, UpdatedAt: base}, + } + got := collectRecentSessions(headers, State{LastRunAt: last}, proj, 20) + if len(got) != 2 { + t.Fatalf("got %d, want 2 (only after LastRunAt)", len(got)) + } + for _, h := range got { + if !h.UpdatedAt.After(last) { + t.Fatalf("session %s not after LastRunAt", h.ID) + } + } +} + +// TestCollectRecentSessionsGlobalOnlyNoMatch: an empty projectDir (global-only +// run) matches nothing, so the window is empty (no-op path). +func TestCollectRecentSessionsGlobalOnlyNoMatch(t *testing.T) { + proj := t.TempDir() + headers := []session.SessionHeader{{ID: "s1", Cwd: proj, UpdatedAt: time.Now().UTC()}} + if got := collectRecentSessions(headers, State{}, "", 20); len(got) != 0 { + t.Fatalf("global-only run must match no sessions, got %d", len(got)) + } +} + +// TestCollectTranscriptsNoSourceOrNoMatch: nil source or no matching session +// yields "" so the caller records Distilled=0 with a "无新增" note. +func TestCollectTranscriptsNoSourceOrNoMatch(t *testing.T) { + if got := collectTranscripts(nil, State{}, t.TempDir(), 20, 0); got != "" { + t.Fatalf("nil source must yield empty transcript, got %q", got) + } + src := &stubSessions{headers: []session.SessionHeader{{ID: "s1", Cwd: t.TempDir(), UpdatedAt: time.Now().UTC()}}} + // Ask for a different (empty) project → no match. + if got := collectTranscripts(src, State{}, "", 20, 0); got != "" { + t.Fatalf("no matching session must yield empty transcript, got %q", got) + } +} + +// TestCollectTranscriptsRendersRoleTagged: matching sessions are rendered into a +// role-tagged transcript containing the session id and message text. +func TestCollectTranscriptsRendersRoleTagged(t *testing.T) { + proj := t.TempDir() + src := &stubSessions{ + headers: []session.SessionHeader{{ID: "sess-abc", Cwd: proj, UpdatedAt: time.Now().UTC()}}, + msgs: map[string]agentcore.MessageList{ + "sess-abc": {userMsg("I always use tabs not spaces"), asstMsg("noted")}, + }, + } + got := collectTranscripts(src, State{}, proj, 20, 0) + for _, want := range []string{"sess-abc", "user: I always use tabs", "assistant: noted"} { + if !strings.Contains(got, want) { + t.Fatalf("transcript missing %q\n---\n%s", want, got) + } + } +} + +// TestParseDistillResponseNewEntries: a canned JSON response yields NewEntry +// writes under the right scope/type dirs, and ephemeral/unknown types are +// dropped. +func TestParseDistillResponseNewEntries(t *testing.T) { + root := "/mem" + proj := t.TempDir() + pid := projectID(proj) + raw := "```json\n" + `{ + "entries": [ + {"type": "user", "scope": "global", "title": "Tabs preference", "body": "Developer prefers tabs over spaces."}, + {"type": "project", "scope": "project", "title": "Arch", "body": "The runner is deterministic; the consolidator is the LLM half."}, + {"type": "todo", "scope": "global", "title": "bad", "body": "ephemeral one-shot task"}, + {"type": "user", "scope": "global", "title": "empty", "body": " "} + ], + "notes": ["distilled 2"] + }` + "\n```" + + entries, notes := parseDistillResponse(raw, nil, root, proj) + if len(entries) != 2 { + t.Fatalf("got %d entries, want 2 (todo + empty dropped): %+v", len(entries), entries) + } + byPath := map[string]string{} + for _, e := range entries { + byPath[e.Path] = e.Body + } + wantGlobal := filepath.Join(root, "global", "user", "tabs-preference.md") + wantProj := filepath.Join(root, "projects", pid, "project", "arch.md") + if _, ok := byPath[wantGlobal]; !ok { + t.Fatalf("missing global user entry at %s; got %v", wantGlobal, byPath) + } + if body, ok := byPath[wantProj]; !ok { + t.Fatalf("missing project entry at %s; got %v", wantProj, byPath) + } else if !strings.HasSuffix(body, "\n") { + t.Fatalf("body must end with newline: %q", body) + } + if strings.Join(notes, "|") == "" { + t.Fatal("expected distill notes surfaced") + } +} + +// TestParseDistillResponseProjectScopeNeedsProjectDir: a "project"-scoped entry on +// a global-only run (empty projectDir) has nowhere valid to live and is skipped. +func TestParseDistillResponseProjectScopeNeedsProjectDir(t *testing.T) { + raw := `{"entries":[{"type":"project","scope":"project","title":"x","body":"y"}]}` + entries, _ := parseDistillResponse(raw, nil, "/mem", "") + if len(entries) != 0 { + t.Fatalf("project entry on global-only run must be skipped, got %+v", entries) + } +} + +// TestParseDistillResponseDedupAgainstExisting: an entry that is a near-duplicate +// of an existing memory is dropped (FR-13). +func TestParseDistillResponseDedupAgainstExisting(t *testing.T) { + existing := []MemoryFile{{Body: "Developer prefers tabs over spaces in all files"}} + raw := `{"entries":[{"type":"user","scope":"global","title":"tabs","body":"Developer prefers tabs over spaces in all files"}]}` + entries, notes := parseDistillResponse(raw, existing, "/mem", "") + if len(entries) != 0 { + t.Fatalf("near-duplicate of existing memory must be dropped, got %+v", entries) + } + if len(notes) == 0 || !strings.Contains(strings.Join(notes, "|"), "near-duplicate") { + t.Fatalf("expected a near-duplicate skip note, got %v", notes) + } +} + +// TestParseDistillResponseUnparseable: an unparseable response adds nothing and +// is not an error (conservative KEEP). +func TestParseDistillResponseUnparseable(t *testing.T) { + entries, notes := parseDistillResponse("the model rambled with no json", nil, "/mem", "") + if entries != nil || notes != nil { + t.Fatalf("unparseable response must yield nothing, got %+v / %v", entries, notes) + } +} + +// TestRunDistillsThroughRunner: an integration test with a stub session source +// and a stub completer-backed llmConsolidator distills a durable fact into a new +// on-disk memory entry, counts it, and reports it. +func TestRunDistillsThroughRunner(t *testing.T) { + root := t.TempDir() + proj := t.TempDir() + + src := &stubSessions{ + headers: []session.SessionHeader{{ID: "s1", Cwd: proj, UpdatedAt: time.Now().UTC()}}, + msgs: map[string]agentcore.MessageList{"s1": {userMsg("always run tests with gotestsum")}}, + } + // llmConsolidator whose distill completion returns one durable fact. The + // merge/prune completion is not reached because there are no eligible files. + cons := &llmConsolidator{complete: func(_ context.Context, _, _ string) (string, error) { + return `{"entries":[{"type":"user","scope":"global","title":"Test runner","body":"Always run tests with gotestsum."}],"notes":["one fact"]}`, nil + }} + r := &Runner{MemoryRoot: root, Consolidator: cons, Sessions: src} + + rep, err := r.Run(context.Background(), RunOptions{ProjectDir: proj}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if rep.Distilled != 1 { + t.Fatalf("Distilled = %d, want 1", rep.Distilled) + } + newPath := filepath.Join(root, "global", "user", "test-runner.md") + raw, err := os.ReadFile(newPath) + if err != nil { + t.Fatalf("distilled entry not written at %s: %v", newPath, err) + } + if !strings.Contains(string(raw), "gotestsum") { + t.Fatalf("distilled body wrong: %q", raw) + } +} + +// TestRunDryRunDistillsButWritesNothing: a dry-run still runs the distill pass +// and reports the count, but writes no new memory files and updates no state. +func TestRunDryRunDistillsButWritesNothing(t *testing.T) { + root := t.TempDir() + proj := t.TempDir() + + src := &stubSessions{ + headers: []session.SessionHeader{{ID: "s1", Cwd: proj, UpdatedAt: time.Now().UTC()}}, + msgs: map[string]agentcore.MessageList{"s1": {userMsg("my stack is Go plus SQLite")}}, + } + cons := &llmConsolidator{complete: func(_ context.Context, _, _ string) (string, error) { + return `{"entries":[{"type":"user","scope":"global","title":"Stack","body":"Stack is Go plus SQLite."}],"notes":[]}`, nil + }} + r := &Runner{MemoryRoot: root, Consolidator: cons, Sessions: src} + + rep, err := r.Run(context.Background(), RunOptions{ProjectDir: proj, DryRun: true}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if rep.Distilled != 1 { + t.Fatalf("dry-run Distilled = %d, want 1 (still reported)", rep.Distilled) + } + if _, err := os.Stat(filepath.Join(root, "global", "user", "stack.md")); !os.IsNotExist(err) { + t.Fatalf("dry-run must not write the distilled entry (err=%v)", err) + } + st, _ := LoadState(root) + if !st.LastRunAt.IsZero() || st.LastStatus != "" { + t.Fatalf("dry-run must not update state: %+v", st) + } +} + +// TestRunNoDurableFactsNoOp: when distillation adds nothing, the report records +// Distilled=0 with the "无新增" note (SPEC §5.5). +func TestRunNoDurableFactsNoOp(t *testing.T) { + root := t.TempDir() + proj := t.TempDir() + src := &stubSessions{ + headers: []session.SessionHeader{{ID: "s1", Cwd: proj, UpdatedAt: time.Now().UTC()}}, + msgs: map[string]agentcore.MessageList{"s1": {userMsg("some transient chatter")}}, + } + cons := &llmConsolidator{complete: func(_ context.Context, _, _ string) (string, error) { + return `{"entries":[],"notes":[]}`, nil + }} + r := &Runner{MemoryRoot: root, Consolidator: cons, Sessions: src} + rep, err := r.Run(context.Background(), RunOptions{ProjectDir: proj}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if rep.Distilled != 0 { + t.Fatalf("Distilled = %d, want 0", rep.Distilled) + } + if !strings.Contains(strings.Join(rep.Notes, "|"), "无新增") { + t.Fatalf("expected '无新增' no-op note, got %v", rep.Notes) + } +} diff --git a/pigo/internal/dream/lock.go b/pigo/internal/dream/lock.go new file mode 100644 index 0000000..c5b5f57 --- /dev/null +++ b/pigo/internal/dream/lock.go @@ -0,0 +1,147 @@ +package dream + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" + "time" +) + +// ErrLocked is returned by AcquireLock when a live (non-stale) lock is already +// held by another process. Callers detect it (via errors.Is) to exit "skipped" +// rather than treating it as a real failure. It is deliberately distinct from +// the I/O errors AcquireLock may also return. +var ErrLocked = errors.New("dream: consolidation already running") + +// DefaultStaleAfter is how long after a lock's started_at the lock is considered +// abandoned (e.g. the holder crashed) and may be taken over. See spec §5.4. +var DefaultStaleAfter = 30 * time.Minute + +// lockInfo is the JSON body persisted in the lock file. +type lockInfo struct { + PID int `json:"pid"` + StartedAt time.Time `json:"started_at"` +} + +// Lock represents an acquired dream single-instance lock backed by the file +// /global/dream/dream.lock. It guarantees at most one consolidation +// runs at a time across processes, with crash-safe stale takeover. +type Lock struct { + path string + released bool +} + +// lockPath is the lock file location under the memory root. It is a separate +// file from state.json and never touches it. +func lockPath(memoryRoot string) string { + return filepath.Join(memoryRoot, "global", "dream", "dream.lock") +} + +// AcquireLock attempts to acquire the dream single-instance lock under +// memoryRoot. On success it returns a *Lock the caller must Release (typically +// via defer). If a live lock is already held it returns ErrLocked. If the +// existing lock is stale (its started_at is older than now-DefaultStaleAfter) or +// malformed/unparseable, it is treated as abandoned and taken over. Any other +// error (permissions, unexpected I/O) is returned as-is so the caller can +// distinguish it from the ErrLocked "skipped" case. +func AcquireLock(memoryRoot string) (*Lock, error) { + dir := filepath.Join(memoryRoot, "global", "dream") + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, err + } + path := lockPath(memoryRoot) + + // First attempt: atomic exclusive create. + l, err := createLock(path) + if err == nil { + return l, nil + } + if !os.IsExist(err) { + // Real I/O error (permissions, etc.), not a contention signal. + return nil, err + } + + // A lock file already exists. Decide whether it is stale and takeable. + if !staleLock(path, time.Now()) { + return nil, ErrLocked + } + + // Stale (or malformed) lock: take it over. Removing then re-creating with + // O_EXCL keeps the create atomic. A racing process that recreates the file + // between our Remove and create will cause our create to fail with EEXIST; + // we surface that as ErrLocked (the other process won the race). + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return nil, err + } + l, err = createLock(path) + if err != nil { + if os.IsExist(err) { + return nil, ErrLocked + } + return nil, err + } + return l, nil +} + +// createLock atomically creates the lock file with O_EXCL and writes the current +// pid + start time as JSON. On EEXIST it returns an error for which os.IsExist +// is true. +func createLock(path string) (*Lock, error) { + f, err := os.OpenFile(path, os.O_EXCL|os.O_CREATE|os.O_WRONLY, 0o644) + if err != nil { + return nil, err + } + data, err := json.Marshal(lockInfo{PID: os.Getpid(), StartedAt: time.Now().UTC()}) + if err != nil { + f.Close() + os.Remove(path) + return nil, err + } + if _, err := f.Write(data); err != nil { + f.Close() + os.Remove(path) + return nil, err + } + if err := f.Close(); err != nil { + os.Remove(path) + return nil, err + } + return &Lock{path: path}, nil +} + +// staleLock reports whether the lock file at path is stale (takeable) as of now. +// A lock is stale when its started_at is older than now-DefaultStaleAfter. A +// missing, unreadable, or malformed/unparseable lock file is also treated as +// stale so a corrupt lock never wedges dream permanently. +func staleLock(path string, now time.Time) bool { + data, err := os.ReadFile(path) + if err != nil { + // Missing or unreadable: treat as takeable. + return true + } + var info lockInfo + if err := json.Unmarshal(data, &info); err != nil { + // Malformed lock body: treat as stale. + return true + } + if info.StartedAt.IsZero() { + // No usable timestamp: treat as stale. + return true + } + return now.Sub(info.StartedAt) > DefaultStaleAfter +} + +// Release removes the lock file. It is safe to call in a defer and safe to +// double-call: a second call (or a call after the file was already removed) is a +// no-op and never panics. A missing file is not an error. +func (l *Lock) Release() error { + if l == nil || l.released { + return nil + } + l.released = true + if err := os.Remove(l.path); err != nil && !os.IsNotExist(err) { + return err + } + return nil +} diff --git a/pigo/internal/dream/lock_test.go b/pigo/internal/dream/lock_test.go new file mode 100644 index 0000000..e8aacbe --- /dev/null +++ b/pigo/internal/dream/lock_test.go @@ -0,0 +1,191 @@ +package dream + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" + "testing" + "time" +) + +// writeLock writes a lock file with the given pid and started_at directly, for +// tests that need to simulate a pre-existing (possibly stale) lock. +func writeLock(t *testing.T, root string, pid int, startedAt time.Time) string { + t.Helper() + dir := filepath.Join(root, "global", "dream") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + path := filepath.Join(dir, "dream.lock") + data, err := json.Marshal(lockInfo{PID: pid, StartedAt: startedAt}) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, data, 0o644); err != nil { + t.Fatal(err) + } + return path +} + +func TestAcquireLockMutualExclusion(t *testing.T) { + root := t.TempDir() + l1, err := AcquireLock(root) + if err != nil { + t.Fatalf("first AcquireLock: %v", err) + } + defer l1.Release() + + // A second acquire while the first is live must fail with ErrLocked. + l2, err := AcquireLock(root) + if !errors.Is(err, ErrLocked) { + t.Fatalf("second AcquireLock err = %v, want ErrLocked", err) + } + if l2 != nil { + t.Errorf("second AcquireLock returned non-nil Lock alongside error") + } +} + +func TestAcquireLockCreatesFileWithBody(t *testing.T) { + root := t.TempDir() + l, err := AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock: %v", err) + } + defer l.Release() + + path := filepath.Join(root, "global", "dream", "dream.lock") + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read lock file: %v", err) + } + var info lockInfo + if err := json.Unmarshal(data, &info); err != nil { + t.Fatalf("lock body not valid JSON: %v", err) + } + if info.PID != os.Getpid() { + t.Errorf("lock pid = %d, want %d", info.PID, os.Getpid()) + } + if info.StartedAt.IsZero() { + t.Errorf("lock started_at is zero, want current time") + } +} + +func TestAcquireLockDoesNotTouchState(t *testing.T) { + root := t.TempDir() + l, err := AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock: %v", err) + } + defer l.Release() + if _, err := os.Stat(filepath.Join(root, "global", "dream", "state.json")); !os.IsNotExist(err) { + t.Errorf("AcquireLock created/left state.json (err=%v); lock must be a separate file", err) + } +} + +func TestAcquireLockStaleTakeover(t *testing.T) { + root := t.TempDir() + // Existing lock older than staleAfter → takeable. + writeLock(t, root, 99999, time.Now().Add(-DefaultStaleAfter-time.Minute)) + + l, err := AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock over stale lock err = %v, want takeover success", err) + } + defer l.Release() + + // The lock body should now reflect our pid. + data, err := os.ReadFile(filepath.Join(root, "global", "dream", "dream.lock")) + if err != nil { + t.Fatal(err) + } + var info lockInfo + if err := json.Unmarshal(data, &info); err != nil { + t.Fatal(err) + } + if info.PID != os.Getpid() { + t.Errorf("after takeover pid = %d, want %d (our pid)", info.PID, os.Getpid()) + } +} + +func TestAcquireLockFreshLockNotTakeable(t *testing.T) { + root := t.TempDir() + // A recently-started lock is live, not stale. + writeLock(t, root, 99999, time.Now().Add(-time.Minute)) + + if _, err := AcquireLock(root); !errors.Is(err, ErrLocked) { + t.Fatalf("AcquireLock over fresh lock err = %v, want ErrLocked", err) + } +} + +func TestAcquireLockMalformedTakeover(t *testing.T) { + root := t.TempDir() + dir := filepath.Join(root, "global", "dream") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "dream.lock"), []byte("{garbage not json"), 0o644); err != nil { + t.Fatal(err) + } + + l, err := AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock over malformed lock err = %v, want takeover success", err) + } + defer l.Release() +} + +func TestReleaseAndReacquire(t *testing.T) { + root := t.TempDir() + l1, err := AcquireLock(root) + if err != nil { + t.Fatalf("first AcquireLock: %v", err) + } + if err := l1.Release(); err != nil { + t.Fatalf("Release: %v", err) + } + // File must be gone after release. + if _, err := os.Stat(filepath.Join(root, "global", "dream", "dream.lock")); !os.IsNotExist(err) { + t.Errorf("lock file still present after Release (err=%v)", err) + } + // Re-acquire must succeed now that the lock is free. + l2, err := AcquireLock(root) + if err != nil { + t.Fatalf("re-acquire after release: %v", err) + } + defer l2.Release() +} + +func TestReleaseDoubleCallSafe(t *testing.T) { + root := t.TempDir() + l, err := AcquireLock(root) + if err != nil { + t.Fatalf("AcquireLock: %v", err) + } + if err := l.Release(); err != nil { + t.Fatalf("first Release: %v", err) + } + // Second Release (and even after the file is gone) must not panic or error. + if err := l.Release(); err != nil { + t.Errorf("second Release err = %v, want nil", err) + } +} + +func TestReleaseNilSafe(t *testing.T) { + var l *Lock + if err := l.Release(); err != nil { + t.Errorf("nil Lock Release err = %v, want nil", err) + } +} + +func TestStaleLockBoundary(t *testing.T) { + root := t.TempDir() + path := writeLock(t, root, 1, time.Unix(0, 0).UTC()) + now := time.Unix(0, 0).UTC() + if staleLock(path, now.Add(DefaultStaleAfter-time.Second)) { + t.Errorf("lock within staleAfter reported stale") + } + if !staleLock(path, now.Add(DefaultStaleAfter+time.Second)) { + t.Errorf("lock past staleAfter reported not stale") + } +} diff --git a/pigo/internal/dream/plan.go b/pigo/internal/dream/plan.go new file mode 100644 index 0000000..16ceaca --- /dev/null +++ b/pigo/internal/dream/plan.go @@ -0,0 +1,359 @@ +package dream + +import ( + "crypto/sha256" + "encoding/hex" + "os" + "path/filepath" + "sort" + "strings" + "unicode" +) + +// NearDupThreshold is the normalized-token Jaccard similarity at or above which +// two memory files are emitted as a near-duplicate candidate pair for the later +// LLM merge-decision step (spec §5.1.1). It is deliberately conservative: this +// node only PAIRS candidates, it never merges, so a false-positive pair costs +// only one extra LLM comparison while a missed pair silently loses a merge +// opportunity. +const NearDupThreshold = 0.7 + +// MemoryFile is one enumerated memory file under the global or the active +// project scope, with the derived metadata the deterministic plan needs. Body +// is retained so the near-duplicate pairing can tokenize it without a second +// read. +type MemoryFile struct { + Path string // absolute path on disk + Scope string // "global" | "projects" + Type string // layout segment ("" for a file directly under the scope root) + Size int64 // byte size of the file content + ContentHash string // sha256 hex of the raw file content + Body string // full file content (used for path extraction + tokenization) +} + +// DedupeGroup is a set of two or more files whose content is byte-identical +// (same ContentHash). The apply node keeps one representative and removes the +// rest; Deduped in the Report counts len(Paths)-1 per group. +type DedupeGroup struct { + Hash string `json:"hash"` + Paths []string `json:"paths"` +} + +// InvalidPathRef is a local filesystem path mentioned in a memory file's body +// that no longer exists on disk. External references (URLs, mailto:, etc.) are +// never recorded here (spec §5.2 / PRD US-004). +type InvalidPathRef struct { + File string `json:"file"` // memory file containing the reference + Ref string `json:"ref"` // the original (unresolved) reference text +} + +// NearDupPair is a candidate pair of files whose token-set similarity is at or +// above NearDupThreshold but whose content is not byte-identical. It is a +// suggestion for the LLM to decide whether an actual merge is warranted; this +// node performs no merge. +type NearDupPair struct { + A string `json:"a"` + B string `json:"b"` + Similarity float64 `json:"similarity"` +} + +// Plan is the plain-data output of the deterministic half of a /dream run. It +// carries no behavior and makes no LLM calls; the later apply node consumes it +// to drive merges/prunes and to compute the final Report counters. +type Plan struct { + Files []MemoryFile `json:"files"` + DedupeGroups []DedupeGroup `json:"dedupe_groups"` + InvalidPathRefs []InvalidPathRef `json:"invalid_path_refs"` + NearDupPairs []NearDupPair `json:"near_dup_pairs"` + BytesBefore int64 `json:"bytes_before"` + FilesBefore int `json:"files_before"` +} + +// BuildPlan enumerates the consolidation-eligible memory files (global scope + +// the given project's scope) under memoryRoot and computes the deterministic +// consolidation plan: exact-dedup groups, invalid local path references, and +// near-duplicate candidate pairs. It excludes the sessions scope entirely and +// any file whose layout type is "checkpoint" (session-transient state, not +// long-term memory — spec §1.3 / §5.1.1). +// +// projectDir is the working directory of the active project; its stable project +// id (matching internal/memory's resolveProjectId) selects the projects +// sub-scope. An empty projectDir restricts the plan to the global scope. A +// missing memoryRoot or scope directory is not an error — it yields an empty +// plan, mirroring memory.Reconcile's tolerance of an absent root. +func BuildPlan(memoryRoot, projectDir string) (Plan, error) { + var plan Plan + + globalRoot := filepath.Join(memoryRoot, "global") + globalFiles, err := enumerateScope(globalRoot, "global") + if err != nil { + return Plan{}, err + } + files := globalFiles + + if projectDir != "" { + projectRoot := filepath.Join(memoryRoot, "projects", projectID(projectDir)) + projFiles, err := enumerateScope(projectRoot, "projects") + if err != nil { + return Plan{}, err + } + files = append(files, projFiles...) + } + + plan.Files = files + plan.FilesBefore = len(files) + for _, f := range files { + plan.BytesBefore += f.Size + } + + plan.DedupeGroups = dedupeGroups(files) + plan.InvalidPathRefs = invalidPathRefs(files, projectDir) + plan.NearDupPairs = nearDupPairs(files) + + return plan, nil +} + +// enumerateScope walks scopeRoot for *.md files, deriving each file's layout +// type from its first path segment relative to scopeRoot. Files under a +// "checkpoint" type directory are skipped. A missing scopeRoot yields no files +// and no error. +func enumerateScope(scopeRoot, scope string) ([]MemoryFile, error) { + var out []MemoryFile + err := filepath.WalkDir(scopeRoot, func(path string, d os.DirEntry, err error) error { + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + if d.IsDir() || !strings.EqualFold(filepath.Ext(d.Name()), ".md") { + return nil + } + rel, relErr := filepath.Rel(scopeRoot, path) + if relErr != nil { + return nil + } + segs := strings.Split(filepath.ToSlash(rel), "/") + typ := "" + if len(segs) >= 2 { + typ = strings.ToLower(segs[0]) + } + if typ == "checkpoint" { + return nil + } + raw, readErr := os.ReadFile(path) + if readErr != nil { + if os.IsNotExist(readErr) { + return nil + } + return readErr + } + sum := sha256.Sum256(raw) + out = append(out, MemoryFile{ + Path: filepath.Clean(path), + Scope: scope, + Type: typ, + Size: int64(len(raw)), + ContentHash: hex.EncodeToString(sum[:]), + Body: string(raw), + }) + return nil + }) + if err != nil { + return nil, err + } + // Deterministic order regardless of directory iteration order. + sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path }) + return out, nil +} + +// dedupeGroups groups files by identical content hash, returning only groups +// with two or more members (a file with a unique hash is not a duplicate). The +// result is stable: groups are ordered by hash, paths within a group by path. +func dedupeGroups(files []MemoryFile) []DedupeGroup { + byHash := make(map[string][]string) + for _, f := range files { + byHash[f.ContentHash] = append(byHash[f.ContentHash], f.Path) + } + var groups []DedupeGroup + for hash, paths := range byHash { + if len(paths) < 2 { + continue + } + sort.Strings(paths) + groups = append(groups, DedupeGroup{Hash: hash, Paths: paths}) + } + sort.Slice(groups, func(i, j int) bool { return groups[i].Hash < groups[j].Hash }) + return groups +} + +// invalidPathRefs scans each file's body for local filesystem path references +// and records the ones that no longer exist. Relative references are resolved +// against projectDir; "~/" against the user home dir. URLs and other external +// references are never considered (see extractLocalPathRefs). +func invalidPathRefs(files []MemoryFile, projectDir string) []InvalidPathRef { + var out []InvalidPathRef + for _, f := range files { + seen := make(map[string]struct{}) + for _, ref := range extractLocalPathRefs(f.Body) { + if _, dup := seen[ref]; dup { + continue + } + seen[ref] = struct{}{} + // A bare relative reference is only meaningful against a project + // base. Without one (global-only plan) skip it rather than + // resolving against an arbitrary cwd and flagging spuriously. + isRelative := !filepath.IsAbs(ref) && !strings.HasPrefix(ref, "~/") + if isRelative && projectDir == "" { + continue + } + resolved := resolveRef(ref, projectDir) + if _, err := os.Stat(resolved); err != nil && os.IsNotExist(err) { + out = append(out, InvalidPathRef{File: f.Path, Ref: ref}) + } + } + } + return out +} + +// nearDupPairs emits candidate pairs whose normalized-token Jaccard similarity +// is at or above NearDupThreshold. Pairs of byte-identical files are skipped: +// those are exact duplicates handled by dedupeGroups, not near-duplicates. +func nearDupPairs(files []MemoryFile) []NearDupPair { + tokens := make([]map[string]struct{}, len(files)) + for i, f := range files { + tokens[i] = tokenize(f.Body) + } + var out []NearDupPair + for i := 0; i < len(files); i++ { + for j := i + 1; j < len(files); j++ { + if files[i].ContentHash == files[j].ContentHash { + continue + } + sim := jaccard(tokens[i], tokens[j]) + if sim >= NearDupThreshold { + out = append(out, NearDupPair{A: files[i].Path, B: files[j].Path, Similarity: sim}) + } + } + } + return out +} + +// projectID derives the stable projects-scope id from a project directory, +// mirroring internal/memory.resolveProjectId: the first 12 hex chars of +// sha256(absPath). The path is made absolute first so relative and absolute +// forms of the same directory map to the same id. +func projectID(projectDir string) string { + abs, err := filepath.Abs(projectDir) + if err != nil { + abs = projectDir + } + sum := sha256.Sum256([]byte(abs)) + return hex.EncodeToString(sum[:])[:12] +} + +// extractLocalPathRefs returns the local filesystem path references found in +// body. It splits on whitespace and common Markdown/code delimiters (so +// `path`, [text](path) and bare tokens are all captured), then keeps only +// tokens that look like a local path while rejecting URLs and other external +// references. +func extractLocalPathRefs(body string) []string { + fields := strings.FieldsFunc(body, func(r rune) bool { + switch r { + case ' ', '\t', '\n', '\r', '`', '(', ')', '[', ']', '{', '}', '"', '\'', '<', '>', ',', ';', '|': + return true + } + return false + }) + var out []string + for _, tok := range fields { + // Trim trailing sentence punctuation that commonly abuts a path. + tok = strings.TrimRight(tok, ".:!?") + if isLocalPathRef(tok) { + out = append(out, tok) + } + } + return out +} + +// isLocalPathRef reports whether tok looks like a local filesystem path rather +// than a URL or other external reference. It accepts absolute ("/..."), +// explicitly relative ("./", "../") and home-relative ("~/") tokens outright; +// a bare relative token (no leading marker) must contain a separator AND a file +// extension in its last segment, so ordinary prose like "TCP/IP", "read/write" +// or "and/or" is not mistaken for a path. URL-schemed tokens are rejected. +func isLocalPathRef(tok string) bool { + if tok == "" { + return false + } + if strings.Contains(tok, "://") { + return false + } + lower := strings.ToLower(tok) + for _, scheme := range []string{"http:", "https:", "ftp:", "ftps:", "file:", "mailto:", "ssh:", "git:", "www."} { + if strings.HasPrefix(lower, scheme) { + return false + } + } + switch { + case strings.HasPrefix(tok, "/"), + strings.HasPrefix(tok, "./"), + strings.HasPrefix(tok, "../"), + strings.HasPrefix(tok, "~/"): + return true + } + // A bare relative token needs both a separator and an extension on its last + // segment to be treated as a path (avoids flagging prose like "TCP/IP"). + if !strings.Contains(tok, "/") { + return false + } + last := tok[strings.LastIndex(tok, "/")+1:] + return strings.Contains(last, ".") && !strings.HasSuffix(last, ".") +} + +// resolveRef turns a reference into an absolute path for existence checking: +// "~/" expands to the user home dir, absolute paths pass through, and relative +// paths resolve against projectDir (or the current dir when projectDir is ""). +func resolveRef(ref, projectDir string) string { + if strings.HasPrefix(ref, "~/") { + if home, err := os.UserHomeDir(); err == nil { + return filepath.Join(home, ref[2:]) + } + } + if filepath.IsAbs(ref) { + return filepath.Clean(ref) + } + return filepath.Join(projectDir, ref) +} + +// tokenize lowercases body and splits it into the set of alphanumeric tokens +// used for near-duplicate similarity. Punctuation and Markdown syntax are +// discarded so formatting differences do not affect the score. +func tokenize(body string) map[string]struct{} { + set := make(map[string]struct{}) + for _, tok := range strings.FieldsFunc(strings.ToLower(body), func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsNumber(r) + }) { + set[tok] = struct{}{} + } + return set +} + +// jaccard returns |a∩b| / |a∪b|. Two empty sets are dissimilar (0) rather than +// identical, so blank files are never paired as near-duplicates. +func jaccard(a, b map[string]struct{}) float64 { + if len(a) == 0 || len(b) == 0 { + return 0 + } + inter := 0 + for t := range a { + if _, ok := b[t]; ok { + inter++ + } + } + union := len(a) + len(b) - inter + if union == 0 { + return 0 + } + return float64(inter) / float64(union) +} diff --git a/pigo/internal/dream/plan_test.go b/pigo/internal/dream/plan_test.go new file mode 100644 index 0000000..29e9de8 --- /dev/null +++ b/pigo/internal/dream/plan_test.go @@ -0,0 +1,216 @@ +package dream + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +// writeMemFile writes body to / creating parent dirs, returning the +// absolute path. +func writeMemFile(t *testing.T, root, rel, body string) string { + t.Helper() + p := filepath.Join(root, rel) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, []byte(body), 0o644); err != nil { + t.Fatal(err) + } + return filepath.Clean(p) +} + +func TestBuildPlanEnumeratesGlobalAndProjectExcludingSessionsAndCheckpoint(t *testing.T) { + root := t.TempDir() + projectDir := t.TempDir() + pid := projectID(projectDir) + + gUser := writeMemFile(t, root, "global/user/prefs.md", "global user prefs\n") + pProj := writeMemFile(t, root, filepath.Join("projects", pid, "project", "arch.md"), "project architecture notes\n") + // Must be excluded: + writeMemFile(t, root, "sessions/sess1/notes/x.md", "session scoped note\n") + writeMemFile(t, root, "global/checkpoint/cp.md", "checkpoint transient\n") + writeMemFile(t, root, filepath.Join("projects", pid, "checkpoint", "cp.md"), "project checkpoint\n") + + plan, err := BuildPlan(root, projectDir) + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + + got := map[string]bool{} + for _, f := range plan.Files { + got[f.Path] = true + } + if !got[gUser] { + t.Errorf("global user file not enumerated") + } + if !got[pProj] { + t.Errorf("project file not enumerated") + } + if plan.FilesBefore != 2 { + t.Errorf("FilesBefore = %d, want 2 (sessions + checkpoint excluded); files=%v", plan.FilesBefore, plan.Files) + } + if plan.BytesBefore == 0 { + t.Errorf("BytesBefore = 0, want >0") + } +} + +func TestBuildPlanEmptyRoot(t *testing.T) { + plan, err := BuildPlan(filepath.Join(t.TempDir(), "does-not-exist"), "") + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + if plan.FilesBefore != 0 || len(plan.Files) != 0 { + t.Errorf("empty root should yield empty plan, got %+v", plan) + } +} + +func TestExactDedupGrouping(t *testing.T) { + root := t.TempDir() + same := "identical memory body\nline two\n" + a := writeMemFile(t, root, "global/user/a.md", same) + b := writeMemFile(t, root, "global/reference/b.md", same) + writeMemFile(t, root, "global/notes/c.md", "a totally different unique body\n") + + plan, err := BuildPlan(root, "") + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + if len(plan.DedupeGroups) != 1 { + t.Fatalf("DedupeGroups = %d, want 1: %+v", len(plan.DedupeGroups), plan.DedupeGroups) + } + g := plan.DedupeGroups[0] + if len(g.Paths) != 2 { + t.Fatalf("group paths = %v, want [a b]", g.Paths) + } + want := map[string]bool{a: true, b: true} + for _, p := range g.Paths { + if !want[p] { + t.Errorf("unexpected path in dedupe group: %q", p) + } + } +} + +func TestPathValidation(t *testing.T) { + root := t.TempDir() + projectDir := t.TempDir() + + // A real file the memory references (relative to projectDir). + existingRel := "src/main.go" + writeMemFile(t, projectDir, existingRel, "package main\n") // reuse helper; writes under projectDir + + body := "See `src/main.go` for the entrypoint.\n" + + "Old helper lived at `src/gone/removed.go` but was deleted.\n" + + "Reference: https://example.com/docs and [site](https://pkg.go.dev/net/http).\n" + + "Email me at mailto:dev@example.com.\n" + writeMemFile(t, root, "global/project/notes.md", body) + + plan, err := BuildPlan(root, projectDir) + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + + var flagged []string + for _, r := range plan.InvalidPathRefs { + flagged = append(flagged, r.Ref) + } + + // Missing local path must be flagged. + if !containsRef(flagged, "src/gone/removed.go") { + t.Errorf("missing path src/gone/removed.go not flagged; flagged=%v", flagged) + } + // Existing local path must NOT be flagged. + if containsRef(flagged, "src/main.go") { + t.Errorf("existing path src/main.go wrongly flagged; flagged=%v", flagged) + } + // URLs / external refs must NEVER be flagged. + for _, r := range flagged { + if wantsURLReject(r) { + t.Errorf("external reference wrongly flagged as invalid local path: %q", r) + } + } +} + +func containsRef(refs []string, want string) bool { + for _, r := range refs { + if r == want { + return true + } + } + return false +} + +func wantsURLReject(r string) bool { + for _, bad := range []string{"http", "https", "mailto", "example.com", "pkg.go.dev"} { + if strings.HasPrefix(r, bad) { + return true + } + } + return false +} + +func TestPathValidationIgnoresProseSlashes(t *testing.T) { + root := t.TempDir() + projectDir := t.TempDir() + // Prose tokens with slashes but no file extension must not be flagged as + // missing local paths. + body := "We support TCP/IP and read/write access; input/output is N/A here.\n" + writeMemFile(t, root, "global/notes/prose.md", body) + + plan, err := BuildPlan(root, projectDir) + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + if len(plan.InvalidPathRefs) != 0 { + t.Errorf("prose slashes wrongly flagged as paths: %+v", plan.InvalidPathRefs) + } +} + +func TestNearDupPairingThreshold(t *testing.T) { + root := t.TempDir() + // Two highly-overlapping (but not identical) bodies -> should pair. + writeMemFile(t, root, "global/user/a.md", + "the quick brown fox jumps over the lazy dog near the river bank today\n") + writeMemFile(t, root, "global/user/b.md", + "the quick brown fox jumps over the lazy dog near the river bank tomorrow\n") + // A dissimilar body -> should not pair with the others. + writeMemFile(t, root, "global/notes/c.md", + "completely unrelated content about database indexing and query planning\n") + + plan, err := BuildPlan(root, "") + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + if len(plan.NearDupPairs) != 1 { + t.Fatalf("NearDupPairs = %d, want exactly 1: %+v", len(plan.NearDupPairs), plan.NearDupPairs) + } + p := plan.NearDupPairs[0] + if p.Similarity < NearDupThreshold { + t.Errorf("paired similarity %v below threshold %v", p.Similarity, NearDupThreshold) + } + // The dissimilar file must not appear in any pair. + for _, pr := range plan.NearDupPairs { + if filepath.Base(pr.A) == "c.md" || filepath.Base(pr.B) == "c.md" { + t.Errorf("dissimilar file c.md wrongly paired: %+v", pr) + } + } +} + +func TestNearDupSkipsExactDuplicates(t *testing.T) { + root := t.TempDir() + same := "one two three four five six seven eight nine ten\n" + writeMemFile(t, root, "global/user/a.md", same) + writeMemFile(t, root, "global/user/b.md", same) + + plan, err := BuildPlan(root, "") + if err != nil { + t.Fatalf("BuildPlan: %v", err) + } + if len(plan.NearDupPairs) != 0 { + t.Errorf("exact duplicates should be handled by dedupe, not near-dup pairs: %+v", plan.NearDupPairs) + } + if len(plan.DedupeGroups) != 1 { + t.Errorf("DedupeGroups = %d, want 1", len(plan.DedupeGroups)) + } +} diff --git a/pigo/internal/dream/prompt.go b/pigo/internal/dream/prompt.go new file mode 100644 index 0000000..a6cbd4b --- /dev/null +++ b/pigo/internal/dream/prompt.go @@ -0,0 +1,90 @@ +package dream + +// This file holds the dream consolidation Agent's system prompt (SPEC §2.2) and +// the strict JSON output schema the model must follow. The prompt scopes the +// task to the LLM half of the mixed division of labor (SPEC §5.1.1): confirm +// semantic merges of near-duplicate entries, and prune only clearly outdated or +// contradicted entries — always conservatively (PRD FR-14: when uncertain, +// KEEP). It never invents facts and only ever names paths that were provided in +// the input, so the deterministic scope guard in the Runner cannot be tricked +// into writing outside the memory store. + +// dreamSystemPrompt is the fixed system instruction for the dream consolidation +// pass. It is deliberately narrow: the Go side already handles exact dedup, path +// validation, near-dup candidate pairing, MEMORY.md rewrite, and index rebuild; +// the model only decides the semantic merges and prunes. +const dreamSystemPrompt = `You are the memory-consolidation agent for pigo ("dream"). You run periodically over a developer's persistent memory library and produce a compact, non-redundant, current set of memory entries. + +Each memory entry is a Markdown file. You are given the current entries (with their absolute file paths and bodies), plus deterministic hints: candidate near-duplicate pairs and references to local files that no longer exist. Exact byte-duplicates and dead-path cleanup are already handled mechanically — you do NOT need to act on those. + +Your job, and ONLY your job: +1. MERGE semantically-overlapping entries. When two or more entries cover the same fact/topic, combine them into a single entry that keeps the most recent and most informative content. Rewrite that surviving entry's full body to be self-contained and concise; the other entries are removed. +2. PRUNE entries that are clearly outdated or directly contradicted by a newer entry. + +Hard rules: +- BE CONSERVATIVE. If you are unsure whether two entries truly overlap, do NOT merge them. If you are unsure whether an entry is outdated or contradicted, KEEP it. Losing a real memory is far worse than leaving a small redundancy. +- NEVER invent facts. A merged body may only restate information already present in the entries you are combining. Do not add, infer, or embellish. +- Only ever reference file paths that appear verbatim in the input. Never emit a path that was not given to you. +- Never merge into, prune, or otherwise target a MEMORY.md index file. Those are indexes, not entries. +- Preserve any Markdown frontmatter (the leading '---' block with name/description/metadata) on a surviving/merged entry, updating it only to reflect the merged content. +- Do NOT create new entries. Distillation of new facts is handled by a separate step. + +Output format: +- Respond with a SINGLE JSON object and nothing else. No prose, no Markdown code fences. +- Schema: + { + "merges": [ + { + "keep": "", + "body": "", + "remove": ["", ...] + } + ], + "prunes": [ + { "path": "", "reason": "" } + ], + "notes": ["", ...] + } +- Every "keep"/"remove"/"path" MUST be one of the input paths. "remove" must not contain the "keep" path. +- If there is nothing to merge or prune, return {"merges": [], "prunes": [], "notes": []}. Returning an empty result is the correct, safe answer when in doubt.` + +// dreamDistillSystemPrompt is the fixed system instruction for the JSONL +// distillation pass (SPEC §5.3, PRD US-005 / FR-13). It is a SEPARATE model call +// from the merge/prune pass above: its input is recent session transcripts plus +// a list of memories that already exist, and its only job is to propose NEW +// durable memory entries that are not already captured. The Go side then dedups +// each proposal against the existing library and path-guards every write, so the +// model only ever supplies type/scope/title/body — never a filesystem path. +const dreamDistillSystemPrompt = `You are the memory-distillation agent for pigo ("dream"). You read recent session transcripts between a developer and an AI coding agent, and extract DURABLE facts worth remembering for future sessions. + +You are also given a list of memories that ALREADY EXIST. Do NOT propose anything already covered by an existing memory — only genuinely new, not-yet-recorded facts. + +What counts as a durable fact (extract these): +- user: stable preferences, conventions, working style, environment the developer states ("I prefer X", "always run tests with Y", "my stack is Z"). +- feedback: corrections or standing instructions the developer gave the agent that should persist. +- project: durable facts about the project's architecture, invariants, key decisions, or layout. +- reference: stable pointers to important resources (a canonical doc, a command, an API) that will remain relevant. + +What to IGNORE (never distill these): +- One-shot task state, TODOs, "now do X" instructions, or anything tied to a single session's in-progress work. +- Ephemeral context: transient errors already fixed, scratch reasoning, temporary file paths. +- Anything you are not confident is durable. When unsure, SKIP it. Recording noise is worse than missing a fact. + +Hard rules: +- BE CONSERVATIVE and specific. Prefer zero entries over speculative ones. Only emit a fact you could justify keeping for months. +- NEVER invent facts. Every entry must be grounded in the transcripts. +- Each entry's body is a short, self-contained Markdown note (a sentence or a few bullet points). Do not include a filesystem path or a filename. +- Classify each entry's "type" as exactly one of: user, feedback, project, reference. +- Classify each entry's "scope" as "project" (specific to the current project) or "global" (applies across all the developer's work). When unsure, use "project". +- Give each entry a short "title" (a few words) used only to name its file. + +Output format: +- Respond with a SINGLE JSON object and nothing else. No prose, no Markdown code fences. +- Schema: + { + "entries": [ + { "type": "user|feedback|project|reference", "scope": "project|global", "title": "", "body": "" } + ], + "notes": ["", ...] + } +- If there is nothing durable to add, return {"entries": [], "notes": []}. An empty result is the correct, safe answer when in doubt.` diff --git a/pigo/internal/dream/reconcile_validation_test.go b/pigo/internal/dream/reconcile_validation_test.go new file mode 100644 index 0000000..f95ec38 --- /dev/null +++ b/pigo/internal/dream/reconcile_validation_test.go @@ -0,0 +1,105 @@ +package dream + +import ( + "context" + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/memory" +) + +// searchHits reopens the memory store that Runner.Run built at +// /index.db and runs a full-text query, returning the set of hit paths. +// The store is reconciled-on-open so the query reflects the exact on-disk state +// left behind by the dream writeback (US-009 / FR-15). The score floor is +// disabled so recall — not ranking — is what the assertion measures. +func searchHits(t *testing.T, root, query string) map[string]bool { + t.Helper() + st, err := memory.Open(filepath.Join(root, "index.db"), root, "") + if err != nil { + t.Fatalf("reopen store: %v", err) + } + defer st.Close() + res, err := st.Search(query, memory.SearchOptions{ScoreFloor: -1, Limit: 50, ReconcileFirst: true}) + if err != nil { + t.Fatalf("Search(%q): %v", query, err) + } + hits := make(map[string]bool, len(res)) + for _, r := range res { + hits[filepath.Clean(r.Path)] = true + } + return hits +} + +// TestReconcileConvergesAfterConsolidation is the US-009 / FR-15 acceptance +// test: after a dream writeback that merges some entries and prunes others, +// memory.Reconcile must rebuild the FTS index and memory_search must converge on +// the compacted current state — never returning the merged-away or pruned +// fragments, always returning the compacted entry. +// +// The scenario uses unique, made-up tokens per fragment so BM25 recall is +// unambiguous: each token exists in exactly one fragment before the run, and the +// assertions check that stale tokens vanish from the index while the compacted +// token appears. +func TestReconcileConvergesAfterConsolidation(t *testing.T) { + root := t.TempDir() + + // Seed distinct (non-duplicate, path-ref-free) fragments so the deterministic + // dedupe/path-clean passes are no-ops and the Consolidator drives the change. + fragA := writeMemFile(t, root, "global/project/frag_a.md", "zorptholine legacy architecture fragment") + fragB := writeMemFile(t, root, "global/project/frag_b.md", "wibblequux duplicate architecture note") + fragC := writeMemFile(t, root, "global/project/frag_c.md", "frobnitz stale prunable outdated entry") + keep := writeMemFile(t, root, "global/user/keep.md", "unrelated grocery shopping list") + + // Stub Consolidator: rewrite frag_a into the compacted current state, merge + // frag_b away into it (deletion), and prune the stale frag_c (deletion). + stub := &stubConsolidator{result: ConsolidateResult{ + MergedBodies: map[string]string{ + fragA: "quombalter consolidated current architecture state", + }, + Deletions: []string{fragB, fragC}, + Merged: 1, + Pruned: 1, + }} + + r := &Runner{MemoryRoot: root, Consolidator: stub} + rep, err := r.Run(context.Background(), RunOptions{}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if !stub.called { + t.Fatal("Consolidator was not called") + } + + // (1) memory.Reconcile ran during writeback and indexed the surviving files. + // The store's index.db is created fresh inside Run, so every kept file is a + // new index row: Indexed must cover the compacted frag_a and the untouched + // keep.md (>=2), proving the FTS index was rebuilt. + if rep.Reconciled.Indexed < 2 { + t.Fatalf("Reconciled.Indexed = %d, want >= 2 (rebuilt index over surviving files)", rep.Reconciled.Indexed) + } + if rep.Merged != 1 || rep.Pruned != 1 { + t.Fatalf("counters not surfaced: Merged=%d Pruned=%d, want 1/1", rep.Merged, rep.Pruned) + } + + // (2) Stale fragments must be gone from the index: neither the merged-away + // fragment (frag_b), the pruned fragment (frag_c), nor the overwritten body of + // frag_a ("zorptholine") may still be searchable. + for _, token := range []string{"zorptholine", "wibblequux", "frobnitz"} { + if hits := searchHits(t, root, token); len(hits) != 0 { + t.Fatalf("stale fragment token %q still searchable after consolidation: %v", token, hits) + } + } + + // (3) The compacted current state IS searchable and resolves to the surviving + // consolidated file (frag_a rewritten in place). + hits := searchHits(t, root, "quombalter") + if !hits[fragA] { + t.Fatalf("compacted entry %q not returned by memory_search for its token, got %v", fragA, hits) + } + + // The unrelated memory must be untouched and still indexed. + if hits := searchHits(t, root, "grocery"); !hits[keep] { + t.Fatalf("untouched memory %q missing from index after consolidation, got %v", keep, hits) + } +} diff --git a/pigo/internal/dream/report.go b/pigo/internal/dream/report.go new file mode 100644 index 0000000..c10a5ec --- /dev/null +++ b/pigo/internal/dream/report.go @@ -0,0 +1,28 @@ +package dream + +// Report is the structured change summary produced by a /dream consolidation +// run. It is persisted inside State.LastReport (state.go) and emitted as a +// single line of JSON on the child process stdout (spec §4.2). Every count is a +// deterministic tally the runner fills in after the plan (Go-deterministic) and +// apply (LLM) phases; the zero value is a valid "nothing changed" report. +// +// The JSON tags match spec §3.2 exactly so the parent process (and any +// scripted/headless caller) can decode the stdout contract without a shared Go +// type. +type Report struct { + Merged int `json:"merged"` // entries merged away by the LLM apply step + Deduped int `json:"deduped"` // exact (content-hash) duplicates removed + PathsCleaned int `json:"paths_cleaned"` // stale local path references cleaned + Pruned int `json:"pruned"` // stale/contradictory entries pruned + Distilled int `json:"distilled"` // new memories distilled from session JSONL + BytesBefore int64 `json:"bytes_before"` + BytesAfter int64 `json:"bytes_after"` + FilesBefore int `json:"files_before"` + FilesAfter int `json:"files_after"` + DryRun bool `json:"dry_run"` + Notes []string `json:"notes,omitempty"` // human-readable reasons (prune causes, etc.) + Reconciled struct { + Indexed int `json:"indexed"` + Pruned int `json:"pruned"` + } `json:"reconciled"` +} diff --git a/pigo/internal/dream/report_test.go b/pigo/internal/dream/report_test.go new file mode 100644 index 0000000..97ac346 --- /dev/null +++ b/pigo/internal/dream/report_test.go @@ -0,0 +1,66 @@ +package dream + +import ( + "encoding/json" + "testing" +) + +func TestReportJSONTags(t *testing.T) { + r := Report{ + Merged: 1, + Deduped: 2, + PathsCleaned: 3, + Pruned: 4, + Distilled: 5, + BytesBefore: 100, + BytesAfter: 80, + FilesBefore: 10, + FilesAfter: 9, + DryRun: true, + Notes: []string{"pruned stale entry"}, + } + r.Reconciled.Indexed = 6 + r.Reconciled.Pruned = 7 + + data, err := json.Marshal(r) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var m map[string]any + if err := json.Unmarshal(data, &m); err != nil { + t.Fatalf("unmarshal: %v", err) + } + for _, key := range []string{ + "merged", "deduped", "paths_cleaned", "pruned", "distilled", + "bytes_before", "bytes_after", "files_before", "files_after", + "dry_run", "notes", "reconciled", + } { + if _, ok := m[key]; !ok { + t.Errorf("missing JSON key %q in %s", key, data) + } + } + rec, ok := m["reconciled"].(map[string]any) + if !ok { + t.Fatalf("reconciled not an object: %v", m["reconciled"]) + } + if _, ok := rec["indexed"]; !ok { + t.Errorf("reconciled.indexed missing: %v", rec) + } + if _, ok := rec["pruned"]; !ok { + t.Errorf("reconciled.pruned missing: %v", rec) + } +} + +func TestReportZeroValueOmitsNotes(t *testing.T) { + data, err := json.Marshal(Report{}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var m map[string]any + if err := json.Unmarshal(data, &m); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if _, ok := m["notes"]; ok { + t.Errorf("empty notes should be omitted, got %s", data) + } +} diff --git a/pigo/internal/dream/runner.go b/pigo/internal/dream/runner.go new file mode 100644 index 0000000..4e24175 --- /dev/null +++ b/pigo/internal/dream/runner.go @@ -0,0 +1,527 @@ +package dream + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/smallnest/pigo/internal/memory" +) + +// Consolidator is the LLM-driven apply step, injected so the deterministic +// Runner skeleton can run end-to-end without an LLM. It receives the plain-data +// plan (dedupe groups + invalid-path refs + near-dup candidate pairs) and +// returns the semantic decisions: which merged bodies to rewrite, which entries +// to delete (merged-away / pruned), and which new entries to distill. +// +// #523 will implement a real Consolidator backed by the main-session model and +// the dream system prompt (SPEC §5.1 step 5, §5.1.1). Until then the Runner uses +// nopConsolidator, so the deterministic half (exact dedup + path clean + +// Reconcile) is fully exercised and testable in isolation. +type Consolidator interface { + Consolidate(ctx context.Context, in ConsolidateInput) (ConsolidateResult, error) +} + +// ConsolidateInput is the plain-data view handed to the Consolidator. It carries +// the deterministic Plan (which already embeds the dedupe groups, invalid-path +// refs and near-dup pairs) plus the resolved roots so the implementation can +// compute in-scope write targets. It intentionally carries no behavior and no +// live handles (no *memory.Store, no *sql.DB) so it stays trivially serializable +// if #523 chooses to marshal it across a subprocess/RPC boundary. +type ConsolidateInput struct { + Plan Plan `json:"plan"` + MemoryRoot string `json:"memory_root"` + ProjectDir string `json:"project_dir"` + // Transcripts is the collected, budget-truncated text of the current + // project's recent session JSONL, gathered deterministically by the Runner + // (SPEC §5.3). It is the input to the separate distillation pass; an empty + // string means there is nothing to distill (no matching sessions), so the + // Consolidator skips the distill call entirely (SPEC §5.5 no-op). + Transcripts string `json:"transcripts,omitempty"` +} + +// NewEntry is a distilled memory file the Consolidator wants created. Path must +// resolve within the memory root's global/project scope; the Runner rejects any +// out-of-scope target before writing (SPEC §5.2 / §7.1). +type NewEntry struct { + Path string `json:"path"` + Body string `json:"body"` +} + +// ConsolidateResult is the Consolidator's decision set. All paths are absolute +// and must lie within the memory root scope. Merged/Pruned/Distilled are the +// report counters the Runner surfaces verbatim; the deterministic Deduped and +// PathsCleaned counters are computed by the Runner itself, not here. +type ConsolidateResult struct { + // MergedBodies maps an existing memory file path to its rewritten (merged / + // compacted) body. The Runner overwrites each file in place. + MergedBodies map[string]string `json:"merged_bodies,omitempty"` + // Deletions are memory files to remove (entries merged away or pruned as + // stale/contradictory). + Deletions []string `json:"deletions,omitempty"` + // NewEntries are freshly distilled memory files to create. + NewEntries []NewEntry `json:"new_entries,omitempty"` + + Merged int `json:"merged"` + Pruned int `json:"pruned"` + Distilled int `json:"distilled"` + Notes []string `json:"notes,omitempty"` +} + +// nopConsolidator is the default no-op Consolidator used when none is injected. +// It makes no decisions, so a Runner with it performs only the deterministic +// dedup + path-clean + Reconcile pass — enough for the skeleton to run +// end-to-end and be unit-tested without an LLM. +type nopConsolidator struct{} + +func (nopConsolidator) Consolidate(context.Context, ConsolidateInput) (ConsolidateResult, error) { + return ConsolidateResult{}, nil +} + +// Runner is the subprocess-side consolidation entry point (SPEC §2.2 / §5.1). It +// runs the deterministic plan, delegates semantic merge/prune/distill to an +// injected Consolidator, applies the results within the memory-root scope, and +// rebuilds the FTS index via memory.Reconcile. The zero value is usable: an +// empty MemoryRoot is resolved from the environment and a nil Consolidator falls +// back to nopConsolidator. +type Runner struct { + // Consolidator is the injected LLM apply step; nil selects nopConsolidator. + Consolidator Consolidator + // MemoryRoot overrides the environment-resolved memory root. Empty means + // resolve via ResolveMemoryRoot (PIGO_HOME / ~/.pigo). Tests set it to a temp + // dir; production leaves it empty. + MemoryRoot string + // Sessions is the source of recent session transcripts for the distillation + // pass (SPEC §5.3). nil resolves the default store at $PIGO_HOME/sessions (or + // ~/.pigo/sessions); tests inject a stub. If it cannot be resolved, + // distillation degrades to a no-op rather than failing the run. + Sessions SessionSource +} + +// RunOptions are the per-invocation parameters mirroring the CLI flags. ProjectDir +// selects the projects sub-scope (empty → global-only); DryRun analyzes without +// writing files or updating state (but still takes the lock — SPEC §5.5). +type RunOptions struct { + DryRun bool + ProjectDir string + // RecentSessions is the first-run distillation window: when dream has never + // run, the most-recent RecentSessions project sessions are distilled (SPEC + // §5.3). Non-positive falls back to DefaultRecentSessions (20). + RecentSessions int +} + +// Run executes one consolidation pass and returns the change Report. The flow is +// the deterministic algorithm of SPEC §5.1: +// +// open store → resolve memoryRoot → acquire lock (ErrLocked → skipped, no error) +// → BuildPlan → Consolidate → if !DryRun: apply dedup + path-clean + consolidation +// → Reconcile → recount → SaveState(ok); if DryRun: counts only, no writes/state. +// +// A held lock is not a failure: Run returns a zero-count Report and a nil error +// so the caller exits 0 ("skipped"). Genuine errors (I/O, plan, apply) are +// returned for the caller to map to exit 1 / status "failed". +func (r *Runner) Run(ctx context.Context, opts RunOptions) (Report, error) { + memoryRoot := r.MemoryRoot + if memoryRoot == "" { + memoryRoot = ResolveMemoryRoot() + } + if memoryRoot == "" { + return Report{}, fmt.Errorf("dream: cannot resolve memory root") + } + + lock, err := AcquireLock(memoryRoot) + if err != nil { + if isLocked(err) { + // Another dream is running: skip silently. Zero-count report, no error + // → caller exits 0, leaves last_status unchanged (SPEC §5.5 / §6.1). We + // have opened nothing and created no files, honoring the skip contract. + return Report{DryRun: opts.DryRun}, nil + } + return Report{}, fmt.Errorf("dream: acquire lock: %w", err) + } + defer lock.Release() + + plan, err := BuildPlan(memoryRoot, opts.ProjectDir) + if err != nil { + return Report{}, fmt.Errorf("dream: build plan: %w", err) + } + + rep := Report{ + DryRun: opts.DryRun, + BytesBefore: plan.BytesBefore, + FilesBefore: plan.FilesBefore, + } + + cons := r.Consolidator + if cons == nil { + cons = nopConsolidator{} + } + + // Gather the current project's recent session transcripts for the distill + // pass (SPEC §5.3). This is deterministic Runner work, mirroring BuildPlan: + // the Consolidator runs the semantic distill call over these transcripts. A + // nil/unresolvable source or no matching session yields "" → the Consolidator + // skips distillation and Distilled stays 0 (SPEC §5.5 no-op). + src := r.Sessions + if src == nil { + if resolved, rerr := resolveSessionStore(); rerr == nil { + src = resolved + } + } + state, _ := LoadState(memoryRoot) + transcripts := collectTranscripts(src, state, opts.ProjectDir, opts.RecentSessions, defaultTranscriptBudget) + + cres, err := cons.Consolidate(ctx, ConsolidateInput{ + Plan: plan, + MemoryRoot: memoryRoot, + ProjectDir: opts.ProjectDir, + Transcripts: transcripts, + }) + if err != nil { + // The runner surfaces the error; the parent/scheduler (node #8) maps a + // non-zero exit to state.LastStatus="failed" (SPEC §4.2/§6.1). The runner + // itself only ever persists "ok", so it never opens/creates the store on a + // failing or dry-run path. + return Report{}, fmt.Errorf("dream: consolidate: %w", err) + } + + // Surface the Consolidator's semantic counters verbatim (SPEC §5.1.1: merge / + // prune / distill are LLM decisions). + rep.Merged = cres.Merged + rep.Pruned = cres.Pruned + rep.Distilled = cres.Distilled + rep.Notes = append(rep.Notes, cres.Notes...) + + // Distillation no-op: no durable facts were added (no matching sessions or + // nothing worth keeping). Record the "无新增" note so the report reflects the + // step ran with no additions (SPEC §5.5, PRD FR-13). + if rep.Distilled == 0 { + rep.Notes = append(rep.Notes, "distill: 无新增") + } + + if opts.DryRun { + // Predict the deterministic counters without touching disk or state, and + // without opening the store (which would create index.db). The after-sizes + // stay zero: nothing was written, so there is no post-state to measure + // (SPEC §5.5 dry-run row). + rep.Deduped = plannedDedupeCount(plan) + rep.PathsCleaned = plannedPathCleanCount(plan) + return rep, nil + } + + // --- write path (non dry-run) ------------------------------------------- + // + // Open the store only here: a dry-run or a lock skip must not create index.db + // (SPEC §5.5). The store is needed solely for the post-write Reconcile. + store, err := memory.Open(filepath.Join(memoryRoot, "index.db"), memoryRoot, "") + if err != nil { + return Report{}, fmt.Errorf("dream: open memory store: %w", err) + } + defer store.Close() + + deleted, deduped, err := applyDedupe(memoryRoot, opts.ProjectDir, plan) + if err != nil { + return Report{}, fmt.Errorf("dream: apply dedupe: %w", err) + } + rep.Deduped = deduped + + cleaned, err := applyPathClean(memoryRoot, opts.ProjectDir, plan, deleted) + if err != nil { + return Report{}, fmt.Errorf("dream: apply path-clean: %w", err) + } + rep.PathsCleaned = cleaned + + if err := applyConsolidation(memoryRoot, opts.ProjectDir, cres, deleted); err != nil { + return Report{}, fmt.Errorf("dream: apply consolidation: %w", err) + } + + // Keep each affected scope's MEMORY.md index consistent with the entries on + // disk: drop any link to a file removed by dedupe or by the Consolidator so + // no dangling references survive (PRD US-003). The full removed set is the + // deterministic dedupe deletions plus the Consolidator's own deletions + // (merged-away + pruned entries). + for _, p := range cres.Deletions { + deleted[filepath.Clean(p)] = struct{}{} + } + if err := updateScopeIndexes(memoryRoot, opts.ProjectDir, deleted); err != nil { + return Report{}, fmt.Errorf("dream: update scope index: %w", err) + } + + res, err := store.Reconcile() + if err != nil { + return Report{}, fmt.Errorf("dream: reconcile: %w", err) + } + rep.Reconciled.Indexed = res.Indexed + rep.Reconciled.Pruned = res.Pruned + + // Recompute post-state sizes by re-enumerating the same scopes. + after, err := BuildPlan(memoryRoot, opts.ProjectDir) + if err != nil { + return Report{}, fmt.Errorf("dream: recount: %w", err) + } + rep.BytesAfter = after.BytesBefore + rep.FilesAfter = after.FilesBefore + + repCopy := rep + if err := SaveState(memoryRoot, State{ + LastRunAt: time.Now().UTC(), + LastStatus: "ok", + LastReport: &repCopy, + }); err != nil { + return Report{}, fmt.Errorf("dream: save state: %w", err) + } + + return rep, nil +} + +// isLocked reports whether err is the ErrLocked contention signal. Kept as a +// helper so the skipped-vs-failed branch reads clearly. +func isLocked(err error) bool { + return errors.Is(err, ErrLocked) +} + +// ResolveMemoryRoot returns the persistent memory root directory the same way +// the CLI does (internal/cli/run.MemoryDir): $PIGO_HOME/memory, else +// ~/.pigo/memory. It is duplicated here rather than imported to keep the dream +// package free of a dependency on the CLI assembly layer (and any import cycle +// through it). Returns "" when neither PIGO_HOME nor the home dir is resolvable. +func ResolveMemoryRoot() string { + dir := os.Getenv("PIGO_HOME") + if dir == "" { + home, err := os.UserHomeDir() + if err != nil { + return "" + } + dir = filepath.Join(home, ".pigo") + } + return filepath.Join(dir, "memory") +} + +// withinScope reports whether target is a write target permitted by this run's +// consolidation scope: it must live under /global or, when a +// projectDir is set, under that project's own /projects/ +// directory. Anything else — the sessions scope, an UNRELATED project's +// directory, or a path outside memoryRoot (e.g. user source) — is rejected. This +// is the SPEC §5.2 / §7.1 path-boundary guard that keeps an LLM-produced path +// from escaping the memory store or clobbering other projects' memories, and it +// mirrors BuildPlan, which only enumerates global + the active project. +// +// The check is symlink-aware: both the scope base and the target's longest +// existing ancestor are passed through filepath.EvalSymlinks before comparison, +// so a symlink planted inside an allowed scope cannot redirect a write/delete +// outside the allowed directory. A prefix match is done on path boundaries (not +// raw string prefixes) so "/globalX" does not pass as "/global". +func withinScope(memoryRoot, projectDir, target string) bool { + if memoryRoot == "" || target == "" { + return false + } + absRoot, err := filepath.Abs(memoryRoot) + if err != nil { + return false + } + absTarget, err := filepath.Abs(target) + if err != nil { + return false + } + resolvedTarget := resolveExisting(filepath.Clean(absTarget)) + + bases := []string{filepath.Join(absRoot, "global")} + if projectDir != "" { + bases = append(bases, filepath.Join(absRoot, "projects", projectID(projectDir))) + } + for _, base := range bases { + base = resolveExisting(base) + if resolvedTarget == base { + return true + } + if strings.HasPrefix(resolvedTarget, base+string(os.PathSeparator)) { + return true + } + } + return false +} + +// resolveExisting returns path with symlinks resolved as far as the filesystem +// allows: it walks up to the longest existing ancestor, resolves that with +// filepath.EvalSymlinks, then rejoins the not-yet-existing tail. This lets the +// scope guard defeat symlink redirection (an existing symlink component is +// dereferenced to its real location) while still handling brand-new target +// paths whose leaf files do not exist yet. On any resolution error it falls back +// to the cleaned input so the guard fails closed via the lexical comparison. +func resolveExisting(path string) string { + path = filepath.Clean(path) + tail := "" + cur := path + for { + if resolved, err := filepath.EvalSymlinks(cur); err == nil { + if tail == "" { + return resolved + } + return filepath.Join(resolved, tail) + } + parent := filepath.Dir(cur) + if parent == cur { + // Reached the root without finding an existing component. + return path + } + tail = filepath.Join(filepath.Base(cur), tail) + cur = parent + } +} + +// plannedDedupeCount is the number of files a dedupe pass would remove: one per +// duplicate beyond the representative in each group (SPEC report.Deduped). +func plannedDedupeCount(plan Plan) int { + n := 0 + for _, g := range plan.DedupeGroups { + if len(g.Paths) > 1 { + n += len(g.Paths) - 1 + } + } + return n +} + +// plannedPathCleanCount is the number of distinct invalid local path references +// a path-clean pass would strip (SPEC report.PathsCleaned). +func plannedPathCleanCount(plan Plan) int { + return len(plan.InvalidPathRefs) +} + +// applyDedupe removes exact-duplicate memory files, keeping the first path in +// each group (paths are pre-sorted by BuildPlan) and deleting the rest. Every +// deletion target is guarded by withinScope. It returns the set of deleted paths +// (so later passes skip them) and the Deduped count. +func applyDedupe(memoryRoot, projectDir string, plan Plan) (map[string]struct{}, int, error) { + deleted := make(map[string]struct{}) + count := 0 + for _, g := range plan.DedupeGroups { + if len(g.Paths) < 2 { + continue + } + // Keep g.Paths[0] as the representative; remove the duplicates. + for _, p := range g.Paths[1:] { + if !withinScope(memoryRoot, projectDir, p) { + return nil, 0, fmt.Errorf("refusing out-of-scope dedupe target %q", p) + } + if err := os.Remove(p); err != nil && !os.IsNotExist(err) { + return nil, 0, err + } + deleted[filepath.Clean(p)] = struct{}{} + count++ + } + } + return deleted, count, nil +} + +// applyPathClean strips invalid local path references from the bodies of the +// files that still exist (skipping any removed by dedupe). Each distinct +// (file, ref) is removed by deleting the exact reference substring and is +// counted once. This is the deterministic half of FR-11; the Consolidator may +// later decide whole-entry pruning for refs that leave an entry meaningless. +// Every rewrite target is guarded by withinScope. +func applyPathClean(memoryRoot, projectDir string, plan Plan, deleted map[string]struct{}) (int, error) { + // Group refs by file so each file is read/written at most once. + byFile := make(map[string][]string) + for _, r := range plan.InvalidPathRefs { + clean := filepath.Clean(r.File) + if _, gone := deleted[clean]; gone { + continue + } + byFile[clean] = append(byFile[clean], r.Ref) + } + + count := 0 + for file, refs := range byFile { + if !withinScope(memoryRoot, projectDir, file) { + return 0, fmt.Errorf("refusing out-of-scope path-clean target %q", file) + } + raw, err := os.ReadFile(file) + if err != nil { + if os.IsNotExist(err) { + continue + } + return 0, err + } + body := string(raw) + for _, ref := range refs { + if strings.Contains(body, ref) { + body = strings.ReplaceAll(body, ref, "") + count++ + } + } + if body != string(raw) { + if err := atomicWrite(file, []byte(body)); err != nil { + return 0, err + } + } + } + return count, nil +} + +// applyConsolidation writes the Consolidator's decisions: rewritten merged +// bodies, new distilled entries, and deletions. Every write/delete target is +// guarded by withinScope so an LLM cannot escape the memory store (SPEC §5.2 / +// §7.1). Deletions already performed by dedupe are skipped. +func applyConsolidation(memoryRoot, projectDir string, cres ConsolidateResult, deleted map[string]struct{}) error { + for path, body := range cres.MergedBodies { + if !withinScope(memoryRoot, projectDir, path) { + return fmt.Errorf("refusing out-of-scope merged-body target %q", path) + } + if err := atomicWrite(path, []byte(body)); err != nil { + return err + } + } + for _, e := range cres.NewEntries { + if !withinScope(memoryRoot, projectDir, e.Path) { + return fmt.Errorf("refusing out-of-scope new-entry target %q", e.Path) + } + if err := atomicWrite(e.Path, []byte(e.Body)); err != nil { + return err + } + } + for _, path := range cres.Deletions { + if _, gone := deleted[filepath.Clean(path)]; gone { + continue + } + if !withinScope(memoryRoot, projectDir, path) { + return fmt.Errorf("refusing out-of-scope deletion target %q", path) + } + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return err + } + } + return nil +} + +// atomicWrite writes data to path via a temp file + rename in the same directory +// so a crash mid-write cannot leave a truncated memory file (SPEC §6.3). The +// parent directory is created lazily. +func atomicWrite(path string, data []byte) error { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + tmp, err := os.CreateTemp(dir, ".dream-*.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + if err := os.Rename(tmpName, path); err != nil { + os.Remove(tmpName) + return err + } + return nil +} diff --git a/pigo/internal/dream/runner_test.go b/pigo/internal/dream/runner_test.go new file mode 100644 index 0000000..c300f28 --- /dev/null +++ b/pigo/internal/dream/runner_test.go @@ -0,0 +1,303 @@ +package dream + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" +) + +// stubConsolidator returns a fixed result and records that it was called. +type stubConsolidator struct { + result ConsolidateResult + called bool +} + +func (s *stubConsolidator) Consolidate(context.Context, ConsolidateInput) (ConsolidateResult, error) { + s.called = true + return s.result, nil +} + +// TestRunEmptyMemoryDir: an empty memory dir yields an all-zero Report with +// status ok (no error). Reconcile tolerates the missing scopes. +func TestRunEmptyMemoryDir(t *testing.T) { + root := t.TempDir() + r := &Runner{MemoryRoot: root} + rep, err := r.Run(context.Background(), RunOptions{}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if rep.FilesBefore != 0 || rep.BytesBefore != 0 || rep.FilesAfter != 0 || rep.BytesAfter != 0 { + t.Fatalf("expected all-zero report, got %+v", rep) + } + if rep.Deduped != 0 || rep.Merged != 0 || rep.Pruned != 0 || rep.PathsCleaned != 0 { + t.Fatalf("expected zero counters, got %+v", rep) + } + // State should be written with status ok for a non-dry-run. + st, err := LoadState(root) + if err != nil { + t.Fatalf("LoadState: %v", err) + } + if st.LastStatus != "ok" { + t.Fatalf("LastStatus = %q, want ok", st.LastStatus) + } +} + +// TestRunDryRunWritesNothing: dry-run computes counts, writes no files, does not +// update state, but still acquires + releases the lock. +func TestRunDryRunWritesNothing(t *testing.T) { + root := t.TempDir() + // Two byte-identical files → one dedupe candidate. + a := writeMemFile(t, root, "global/user/a.md", "same content") + b := writeMemFile(t, root, "global/user/b.md", "same content") + + r := &Runner{MemoryRoot: root} + rep, err := r.Run(context.Background(), RunOptions{DryRun: true}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if !rep.DryRun { + t.Fatal("Report.DryRun = false, want true") + } + if rep.Deduped != 1 { + t.Fatalf("Deduped = %d, want 1 (predicted)", rep.Deduped) + } + // Both files must still exist — dry-run writes nothing. + if _, err := os.Stat(a); err != nil { + t.Fatalf("file a removed in dry-run: %v", err) + } + if _, err := os.Stat(b); err != nil { + t.Fatalf("file b removed in dry-run: %v", err) + } + // State must NOT be updated (never-run remains). + st, err := LoadState(root) + if err != nil { + t.Fatalf("LoadState: %v", err) + } + if !st.LastRunAt.IsZero() || st.LastStatus != "" { + t.Fatalf("dry-run updated state: %+v", st) + } + // Lock must have been released (a fresh acquire succeeds). + lk, err := AcquireLock(root) + if err != nil { + t.Fatalf("lock not released after dry-run: %v", err) + } + lk.Release() +} + +// TestRunLockedSkips: when a live lock is already held, Run returns a zero-count +// report and NO error (exit-0 "skipped" semantics), and does not touch state. +func TestRunLockedSkips(t *testing.T) { + root := t.TempDir() + writeMemFile(t, root, "global/user/a.md", "content") + + held, err := AcquireLock(root) + if err != nil { + t.Fatalf("pre-acquire lock: %v", err) + } + defer held.Release() + + r := &Runner{MemoryRoot: root} + rep, err := r.Run(context.Background(), RunOptions{}) + if err != nil { + t.Fatalf("Run under held lock should not error, got %v", err) + } + if rep.FilesBefore != 0 || rep.Deduped != 0 { + t.Fatalf("skipped run should have zero report, got %+v", rep) + } + // State untouched (never ran). + st, _ := LoadState(root) + if !st.LastRunAt.IsZero() || st.LastStatus != "" { + t.Fatalf("skipped run touched state: %+v", st) + } +} + +// TestRunAppliesDedupe: a non-dry-run removes exact duplicates, updates state, +// and reflects counts in the Report. +func TestRunAppliesDedupe(t *testing.T) { + root := t.TempDir() + a := writeMemFile(t, root, "global/user/a.md", "same content") + b := writeMemFile(t, root, "global/user/b.md", "same content") + + r := &Runner{MemoryRoot: root} + rep, err := r.Run(context.Background(), RunOptions{}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if rep.Deduped != 1 { + t.Fatalf("Deduped = %d, want 1", rep.Deduped) + } + // Exactly one of the two duplicates must survive (the sorted-first: a.md). + _, errA := os.Stat(a) + _, errB := os.Stat(b) + if errA != nil { + t.Fatalf("representative a.md removed: %v", errA) + } + if errB == nil { + t.Fatal("duplicate b.md should have been removed") + } + if rep.FilesAfter != 1 { + t.Fatalf("FilesAfter = %d, want 1", rep.FilesAfter) + } + st, _ := LoadState(root) + if st.LastStatus != "ok" || st.LastReport == nil { + t.Fatalf("state not updated: %+v", st) + } +} + +// TestWithinScope: the path-boundary guard accepts in-scope targets and rejects +// everything outside /global and the active project's directory. +func TestWithinScope(t *testing.T) { + root := t.TempDir() + projectDir := t.TempDir() + pid := projectID(projectDir) + otherPID := "0123456789ab" // a different, unrelated project id + cases := []struct { + name string + project string + target string + want bool + }{ + {"global file", projectDir, filepath.Join(root, "global", "user", "x.md"), true}, + {"active project file", projectDir, filepath.Join(root, "projects", pid, "notes", "y.md"), true}, + {"global root itself", projectDir, filepath.Join(root, "global"), true}, + {"unrelated project rejected", projectDir, filepath.Join(root, "projects", otherPID, "z.md"), false}, + {"any project rejected when global-only", "", filepath.Join(root, "projects", pid, "y.md"), false}, + {"global still ok when global-only", "", filepath.Join(root, "global", "x.md"), true}, + {"sessions scope rejected", projectDir, filepath.Join(root, "sessions", "s1", "checkpoint.md"), false}, + {"outside root rejected", projectDir, filepath.Join(root, "..", "evil.md"), false}, + {"sibling prefix not confused", projectDir, filepath.Join(root, "globalX", "z.md"), false}, + {"absolute escape rejected", projectDir, "/etc/passwd", false}, + {"empty target rejected", projectDir, "", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := withinScope(root, tc.project, tc.target); got != tc.want { + t.Fatalf("withinScope(%q, %q, %q) = %v, want %v", root, tc.project, tc.target, got, tc.want) + } + }) + } + if withinScope("", projectDir, filepath.Join(root, "global", "x.md")) { + t.Fatal("empty memoryRoot must reject") + } +} + +// TestWithinScopeSymlinkEscape: a symlink planted inside an allowed scope must +// not let a target escape memoryRoot. The guard resolves symlinks on existing +// ancestors before the containment check (SPEC §7.1 defense-in-depth). +func TestWithinScopeSymlinkEscape(t *testing.T) { + root := t.TempDir() + outside := t.TempDir() // a directory fully outside memoryRoot + + globalDir := filepath.Join(root, "global") + if err := os.MkdirAll(globalDir, 0o755); err != nil { + t.Fatalf("mkdir global: %v", err) + } + // /global/escape -> + link := filepath.Join(globalDir, "escape") + if err := os.Symlink(outside, link); err != nil { + t.Skipf("symlink unsupported: %v", err) + } + + // Lexically this looks in-scope (/global/escape/evil.md) but resolves + // to /evil.md, which must be rejected. + target := filepath.Join(link, "evil.md") + if withinScope(root, "", target) { + t.Fatalf("symlink escape target accepted: %q", target) + } +} + +// TestRunPathClean: a memory file referencing a non-existent local path has that +// reference stripped and counted. +func TestRunPathClean(t *testing.T) { + root := t.TempDir() + proj := t.TempDir() + missing := filepath.Join(proj, "does", "not", "exist.go") + body := "See `" + missing + "` for details." + f := writeMemFile(t, root, "global/reference/r.md", body) + + r := &Runner{MemoryRoot: root} + rep, err := r.Run(context.Background(), RunOptions{ProjectDir: proj}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if rep.PathsCleaned != 1 { + t.Fatalf("PathsCleaned = %d, want 1", rep.PathsCleaned) + } + raw, err := os.ReadFile(f) + if err != nil { + t.Fatalf("read cleaned file: %v", err) + } + if got := string(raw); got == body { + t.Fatalf("body unchanged after path-clean: %q", got) + } +} + +// TestRunConsolidatorApplied: an injected Consolidator's new-entry write and +// counters flow through, and the write lands within scope. +func TestRunConsolidatorApplied(t *testing.T) { + root := t.TempDir() + writeMemFile(t, root, "global/user/a.md", "hello") + + newPath := filepath.Join(root, "global", "user", "distilled.md") + stub := &stubConsolidator{result: ConsolidateResult{ + NewEntries: []NewEntry{{Path: newPath, Body: "distilled fact"}}, + Distilled: 1, + Merged: 2, + Pruned: 3, + Notes: []string{"note"}, + }} + r := &Runner{MemoryRoot: root, Consolidator: stub} + rep, err := r.Run(context.Background(), RunOptions{}) + if err != nil { + t.Fatalf("Run: %v", err) + } + if !stub.called { + t.Fatal("Consolidator was not called") + } + if rep.Distilled != 1 || rep.Merged != 2 || rep.Pruned != 3 { + t.Fatalf("counters not surfaced: %+v", rep) + } + if _, err := os.Stat(newPath); err != nil { + t.Fatalf("new entry not written: %v", err) + } +} + +// TestApplyConsolidationRejectsOutOfScope: a Consolidator that tries to write +// outside the memory root is rejected by the guard. +func TestApplyConsolidationRejectsOutOfScope(t *testing.T) { + root := t.TempDir() + outside := filepath.Join(t.TempDir(), "escape.md") + cres := ConsolidateResult{ + NewEntries: []NewEntry{{Path: outside, Body: "x"}}, + } + if err := applyConsolidation(root, "", cres, nil); err == nil { + t.Fatal("expected out-of-scope write to be rejected") + } + if _, err := os.Stat(outside); err == nil { + t.Fatal("out-of-scope file must not be created") + } +} + +// TestRunDryRunLeavesStaleLockTakeable is a small guard that dry-run's lock is +// released promptly (no leftover live lock). +func TestRunDryRunLockReleased(t *testing.T) { + root := t.TempDir() + r := &Runner{MemoryRoot: root} + if _, err := r.Run(context.Background(), RunOptions{DryRun: true}); err != nil { + t.Fatalf("Run: %v", err) + } + // Fresh acquire must succeed immediately (lock released, not merely stale). + // Force a long stale window so a leftover *live* lock would block acquisition, + // proving Run released it rather than leaving it to be reclaimed as stale. + // Restore the package default afterward so this does not leak into other tests. + orig := DefaultStaleAfter + DefaultStaleAfter = time.Hour + defer func() { DefaultStaleAfter = orig }() + lk, err := AcquireLock(root) + if err != nil { + t.Fatalf("lock not released: %v", err) + } + lk.Release() +} diff --git a/pigo/internal/dream/scheduler.go b/pigo/internal/dream/scheduler.go new file mode 100644 index 0000000..c3b36db --- /dev/null +++ b/pigo/internal/dream/scheduler.go @@ -0,0 +1,103 @@ +package dream + +import ( + "context" + "time" +) + +// BackgroundSpawn launches one dream consolidation subprocess for projectDir and +// returns its decoded Report. Implementations live in the CLI layer (they shell +// out to `pigo --dream -C ` and parse the stdout Report), so this +// package stays free of os/exec concerns and there is no import cycle back +// through the CLI. A non-nil error means the run failed; background failures are +// silent to the user (the subprocess itself records last_status="failed"), so +// MaybeRunBackground swallows the error rather than surfacing it. +type BackgroundSpawn func(ctx context.Context, projectDir string) (Report, error) + +// BackgroundDeps carries everything MaybeRunBackground needs to decide on and run +// a startup auto-consolidation without pulling process/exec or presentation +// concerns into this package. +type BackgroundDeps struct { + // MemoryRoot is the dream state root read to decide whether a run is due. It + // must match the root the subprocess consolidates (dream.ResolveMemoryRoot), + // so the parent's Due check sees the same last_run_at the child updates. + MemoryRoot string + // ProjectDir is the working directory attributed to the run (project scope). + ProjectDir string + // Config is the resolved [dream] configuration (enabled / interval). + Config Config + // Now supplies the current time for the due check; nil uses time.Now. It is a + // seam so tests can drive Due deterministically. + Now func() time.Time + // Spawn launches the subprocess. When nil, MaybeRunBackground does nothing. + Spawn BackgroundSpawn + // OnReport is invoked (from the background goroutine) with the completed + // report only when the run produced actual changes — worth a one-line notice. + // A skipped run (another dream held the lock → all-zero report) or a no-op run + // yields no call, keeping the trigger non-intrusive. Nil disables the notice. + OnReport func(Report) +} + +// Scheduler owns the startup auto-trigger decision (SPEC §2.1 Scheduler +// component). It is stateless: Due reads state.json on demand and +// MaybeRunBackground spawns at most one background run per call. The +// single-instance guarantee is enforced by the subprocess's O_EXCL lock, not +// here — a second trigger simply results in a skipped child. +type Scheduler struct{} + +// Due reports whether an auto-triggered consolidation is warranted now. It is +// cheap by design (SPEC §8.2 zero-startup-overhead): when dream is disabled it +// returns immediately without touching the filesystem; otherwise it reads +// state.json once and defers to State.Due (which also returns false for a +// never-run zero LastRunAt, so the first-ever run is never auto-triggered). +func (Scheduler) Due(memoryRoot string, cfg Config, now time.Time) bool { + if !cfg.Enabled { + return false + } + st, _ := LoadState(memoryRoot) + return st.Due(cfg, now) +} + +// MaybeRunBackground checks (cheaply) whether a consolidation is due and, if so, +// spawns it in a detached goroutine and returns immediately — it never blocks +// the caller, so the first interactive response is never delayed (SPEC FR-4 / +// §8.2). It returns true when a background run was launched. When dream is +// disabled or not due it returns false after at most a single state.json read +// (no goroutine, no subprocess). +// +// On completion the goroutine surfaces a one-line notice via OnReport only for a +// run that changed something; a skipped run (lock held elsewhere → zero report), +// a no-op run, or a failed run is silent (SPEC §6.1 background row). +func (s Scheduler) MaybeRunBackground(ctx context.Context, deps BackgroundDeps) bool { + if deps.Spawn == nil { + return false + } + now := time.Now + if deps.Now != nil { + now = deps.Now + } + if !s.Due(deps.MemoryRoot, deps.Config, now()) { + return false + } + go func() { + rep, err := deps.Spawn(ctx, deps.ProjectDir) + if err != nil { + // Background failure: silent. The subprocess already recorded + // last_status="failed"; we do not interrupt the user with an error. + return + } + if deps.OnReport != nil && reportHasChanges(rep) { + deps.OnReport(rep) + } + }() + return true +} + +// reportHasChanges reports whether r reflects any actual mutation. A background +// run that skipped (lock contention) or found nothing to do produces an all-zero +// report, which is not worth a startup notice. +func reportHasChanges(r Report) bool { + return r.Merged > 0 || r.Deduped > 0 || r.PathsCleaned > 0 || + r.Pruned > 0 || r.Distilled > 0 || + r.Reconciled.Indexed > 0 || r.Reconciled.Pruned > 0 +} diff --git a/pigo/internal/dream/scheduler_test.go b/pigo/internal/dream/scheduler_test.go new file mode 100644 index 0000000..ef9ff3a --- /dev/null +++ b/pigo/internal/dream/scheduler_test.go @@ -0,0 +1,211 @@ +package dream + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +// fixedNow returns a func yielding t, for the Now seam. +func fixedNow(t time.Time) func() time.Time { return func() time.Time { return t } } + +// writeDueState seeds a state.json under memoryRoot with a LastRunAt old enough +// that Due(cfg, now) is true for the default interval. +func writeDueState(t *testing.T, memoryRoot string, lastRun time.Time) { + t.Helper() + if err := SaveState(memoryRoot, State{LastRunAt: lastRun, LastStatus: "ok"}); err != nil { + t.Fatalf("SaveState: %v", err) + } +} + +func TestSchedulerDue(t *testing.T) { + root := t.TempDir() + now := time.Date(2026, 1, 20, 12, 0, 0, 0, time.UTC) + cfg := NewConfig(nil, 7, 20) // enabled, 7-day interval + + var s Scheduler + + // Disabled → never due, and cheap (no state read needed). + disabled := NewConfig(boolPtr(false), 7, 20) + if s.Due(root, disabled, now) { + t.Fatal("disabled config must not be due") + } + + // No state file (never run) → not due (first run is manual, spec §11.1). + if s.Due(root, cfg, now) { + t.Fatal("never-run state must not be due") + } + + // Last run within the interval → not due. + writeDueState(t, root, now.Add(-3*24*time.Hour)) + if s.Due(root, cfg, now) { + t.Fatal("run 3d ago with 7d interval must not be due") + } + + // Last run older than the interval → due. + writeDueState(t, root, now.Add(-8*24*time.Hour)) + if !s.Due(root, cfg, now) { + t.Fatal("run 8d ago with 7d interval must be due") + } +} + +func TestMaybeRunBackground_NotSpawnedWhenDisabled(t *testing.T) { + root := t.TempDir() + writeDueState(t, root, time.Now().Add(-30*24*time.Hour)) // would be due if enabled + + var spawned bool + launched := Scheduler{}.MaybeRunBackground(context.Background(), BackgroundDeps{ + MemoryRoot: root, + Config: NewConfig(boolPtr(false), 7, 20), + Now: fixedNow(time.Now()), + Spawn: func(context.Context, string) (Report, error) { + spawned = true + return Report{}, nil + }, + }) + if launched { + t.Fatal("MaybeRunBackground returned true for disabled dream") + } + if spawned { + t.Fatal("subprocess must not be spawned when disabled") + } +} + +func TestMaybeRunBackground_NotSpawnedWhenNotDue(t *testing.T) { + root := t.TempDir() + now := time.Date(2026, 2, 1, 9, 0, 0, 0, time.UTC) + writeDueState(t, root, now.Add(-1*24*time.Hour)) // 1 day ago, interval 7d → not due + + var spawned bool + launched := Scheduler{}.MaybeRunBackground(context.Background(), BackgroundDeps{ + MemoryRoot: root, + Config: NewConfig(nil, 7, 20), + Now: fixedNow(now), + Spawn: func(context.Context, string) (Report, error) { + spawned = true + return Report{}, nil + }, + }) + if launched || spawned { + t.Fatalf("not-due run must not launch/spawn (launched=%v spawned=%v)", launched, spawned) + } +} + +func TestMaybeRunBackground_SpawnsAndNoticesOnChanges(t *testing.T) { + root := t.TempDir() + now := time.Date(2026, 2, 10, 9, 0, 0, 0, time.UTC) + writeDueState(t, root, now.Add(-10*24*time.Hour)) // due + + var ( + mu sync.Mutex + gotDir string + reported *Report + done = make(chan struct{}) + ) + launched := Scheduler{}.MaybeRunBackground(context.Background(), BackgroundDeps{ + MemoryRoot: root, + ProjectDir: "/proj/x", + Config: NewConfig(nil, 7, 20), + Now: fixedNow(now), + Spawn: func(_ context.Context, dir string) (Report, error) { + mu.Lock() + gotDir = dir + mu.Unlock() + return Report{Merged: 2, Deduped: 1}, nil + }, + OnReport: func(r Report) { + mu.Lock() + reported = &r + mu.Unlock() + close(done) + }, + }) + if !launched { + t.Fatal("due run must launch a background spawn") + } + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("OnReport not called within timeout") + } + mu.Lock() + defer mu.Unlock() + if gotDir != "/proj/x" { + t.Fatalf("spawn got dir %q, want /proj/x", gotDir) + } + if reported == nil || reported.Merged != 2 { + t.Fatalf("OnReport got %+v, want Merged=2", reported) + } +} + +func TestMaybeRunBackground_SkippedRunIsSilent(t *testing.T) { + root := t.TempDir() + now := time.Date(2026, 3, 1, 9, 0, 0, 0, time.UTC) + writeDueState(t, root, now.Add(-10*24*time.Hour)) // due + + spawnDone := make(chan struct{}) + var noticed bool + var mu sync.Mutex + Scheduler{}.MaybeRunBackground(context.Background(), BackgroundDeps{ + MemoryRoot: root, + Config: NewConfig(nil, 7, 20), + Now: fixedNow(now), + Spawn: func(context.Context, string) (Report, error) { + // A skipped/lock-held run emits an all-zero report with no error. + defer close(spawnDone) + return Report{}, nil + }, + OnReport: func(Report) { + mu.Lock() + noticed = true + mu.Unlock() + }, + }) + <-spawnDone + // Give the goroutine a moment past Spawn to (not) call OnReport. + time.Sleep(20 * time.Millisecond) + mu.Lock() + defer mu.Unlock() + if noticed { + t.Fatal("all-zero (skipped/no-op) report must not produce a notice") + } +} + +func TestMaybeRunBackground_FailureIsSilent(t *testing.T) { + root := t.TempDir() + now := time.Date(2026, 3, 5, 9, 0, 0, 0, time.UTC) + writeDueState(t, root, now.Add(-10*24*time.Hour)) // due + + spawnDone := make(chan struct{}) + var noticed bool + var mu sync.Mutex + Scheduler{}.MaybeRunBackground(context.Background(), BackgroundDeps{ + MemoryRoot: root, + Config: NewConfig(nil, 7, 20), + Now: fixedNow(now), + Spawn: func(context.Context, string) (Report, error) { + defer close(spawnDone) + return Report{Merged: 5}, errors.New("boom") + }, + OnReport: func(Report) { + mu.Lock() + noticed = true + mu.Unlock() + }, + }) + <-spawnDone + time.Sleep(20 * time.Millisecond) + mu.Lock() + defer mu.Unlock() + if noticed { + t.Fatal("a failed background run must be silent (no notice)") + } +} + +func TestMaybeRunBackground_NilSpawn(t *testing.T) { + if (Scheduler{}).MaybeRunBackground(context.Background(), BackgroundDeps{}) { + t.Fatal("nil Spawn must yield false (no-op)") + } +} diff --git a/pigo/internal/dream/state.go b/pigo/internal/dream/state.go new file mode 100644 index 0000000..2242143 --- /dev/null +++ b/pigo/internal/dream/state.go @@ -0,0 +1,99 @@ +package dream + +import ( + "encoding/json" + "os" + "path/filepath" + "time" +) + +// State is the persisted /dream run state, stored as JSON at +// /global/dream/state.json. A zero-value State (LastRunAt zero, +// empty status, nil report) means dream has never run. +type State struct { + LastRunAt time.Time `json:"last_run_at"` + LastStatus string `json:"last_status"` // "ok" | "failed" | "skipped" + // LastReport holds the structured change report from the last run. It is + // nil until dream has completed at least one non-dry-run pass. + LastReport *Report `json:"last_report,omitempty"` +} + +// statePath is the state file location under the memory root. +func statePath(memoryRoot string) string { + return filepath.Join(memoryRoot, "global", "dream", "state.json") +} + +// LoadState reads the dream state from /global/dream/state.json. A +// missing file returns a zero-value State (never run) with no error. Corrupt or +// unreadable JSON is tolerated the same way: the caller gets a zero-value State +// and no error, so a damaged state file degrades to "never run" rather than +// breaking dream entirely. +func LoadState(memoryRoot string) (State, error) { + data, err := os.ReadFile(statePath(memoryRoot)) + if err != nil { + if os.IsNotExist(err) { + return State{}, nil + } + // Unreadable (permissions, transient IO): treat as never-run rather + // than surfacing an error that would block dream. + return State{}, nil + } + var s State + if err := json.Unmarshal(data, &s); err != nil { + // Corrupt JSON: degrade to never-run. + return State{}, nil + } + return s, nil +} + +// SaveState writes the dream state to /global/dream/state.json, +// creating the parent directory lazily. The file is written atomically via a +// temp file + rename so a crash mid-write cannot leave a truncated state.json. +func SaveState(memoryRoot string, s State) error { + dir := filepath.Join(memoryRoot, "global", "dream") + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + data, err := json.MarshalIndent(s, "", " ") + if err != nil { + return err + } + path := statePath(memoryRoot) + tmp, err := os.CreateTemp(dir, "state-*.json.tmp") + if err != nil { + return err + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return err + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return err + } + if err := os.Rename(tmpName, path); err != nil { + os.Remove(tmpName) + return err + } + return nil +} + +// Due reports whether an auto-triggered consolidation is warranted at now given +// cfg. It returns true only when dream is enabled and the configured interval +// has elapsed since the last run. Two cases deliberately return false: +// - cfg.Enabled is false: auto-trigger is disabled entirely (US-001). +// - a zero LastRunAt (dream has never run): the first-ever run is NOT +// auto-triggered — the user is prompted to run /dream manually instead — to +// avoid a cold-start token cost for new users. See spec §11.1. +func (s State) Due(cfg Config, now time.Time) bool { + if !cfg.Enabled { + return false + } + if s.LastRunAt.IsZero() { + return false + } + interval := time.Duration(cfg.IntervalDays) * 24 * time.Hour + return now.Sub(s.LastRunAt) >= interval +} diff --git a/pigo/internal/dream/state_test.go b/pigo/internal/dream/state_test.go new file mode 100644 index 0000000..b29208e --- /dev/null +++ b/pigo/internal/dream/state_test.go @@ -0,0 +1,100 @@ +package dream + +import ( + "os" + "path/filepath" + "testing" + "time" +) + +func TestStateRoundTrip(t *testing.T) { + root := t.TempDir() + report := &Report{Merged: 3, Deduped: 1, DryRun: false} + want := State{ + LastRunAt: time.Date(2026, 8, 1, 9, 0, 0, 0, time.UTC), + LastStatus: "ok", + LastReport: report, + } + if err := SaveState(root, want); err != nil { + t.Fatalf("SaveState: %v", err) + } + got, err := LoadState(root) + if err != nil { + t.Fatalf("LoadState: %v", err) + } + if !got.LastRunAt.Equal(want.LastRunAt) { + t.Errorf("LastRunAt = %v, want %v", got.LastRunAt, want.LastRunAt) + } + if got.LastStatus != want.LastStatus { + t.Errorf("LastStatus = %q, want %q", got.LastStatus, want.LastStatus) + } + if got.LastReport == nil { + t.Fatalf("LastReport = nil, want %+v", want.LastReport) + } + if got.LastReport.Merged != report.Merged || got.LastReport.Deduped != report.Deduped { + t.Errorf("LastReport = %+v, want %+v", got.LastReport, report) + } +} + +func TestSaveStateCreatesDir(t *testing.T) { + root := t.TempDir() + if err := SaveState(root, State{LastStatus: "ok"}); err != nil { + t.Fatalf("SaveState: %v", err) + } + if _, err := os.Stat(filepath.Join(root, "global", "dream", "state.json")); err != nil { + t.Errorf("state.json not created: %v", err) + } +} + +func TestLoadStateMissingIsZero(t *testing.T) { + got, err := LoadState(t.TempDir()) + if err != nil { + t.Fatalf("LoadState: %v", err) + } + if !got.LastRunAt.IsZero() || got.LastStatus != "" || got.LastReport != nil { + t.Errorf("missing state = %+v, want zero-value", got) + } +} + +func TestLoadStateCorruptTolerated(t *testing.T) { + root := t.TempDir() + dir := filepath.Join(root, "global", "dream") + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "state.json"), []byte("{not valid json"), 0o644); err != nil { + t.Fatal(err) + } + got, err := LoadState(root) + if err != nil { + t.Fatalf("LoadState returned error on corrupt JSON, want tolerated: %v", err) + } + if !got.LastRunAt.IsZero() { + t.Errorf("corrupt state = %+v, want zero-value (never run)", got) + } +} + +func TestDue(t *testing.T) { + now := time.Date(2026, 8, 1, 12, 0, 0, 0, time.UTC) + cfg := Config{Enabled: true, IntervalDays: 7, RecentSessions: 20} + tests := []struct { + name string + cfg Config + last time.Time + want bool + }{ + {"due: interval elapsed", cfg, now.Add(-8 * 24 * time.Hour), true}, + {"due: exactly at interval", cfg, now.Add(-7 * 24 * time.Hour), true}, + {"not due: within interval", cfg, now.Add(-3 * 24 * time.Hour), false}, + {"zero LastRunAt never due", cfg, time.Time{}, false}, + {"disabled never due", Config{Enabled: false, IntervalDays: 7}, now.Add(-30 * 24 * time.Hour), false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := State{LastRunAt: tt.last} + if got := s.Due(tt.cfg, now); got != tt.want { + t.Errorf("Due = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/pigo/internal/hooks/config.go b/pigo/internal/hooks/config.go new file mode 100644 index 0000000..c090fd5 --- /dev/null +++ b/pigo/internal/hooks/config.go @@ -0,0 +1,83 @@ +// Package hooks implements pigo's user-extensible lifecycle hook system: a +// config-driven way to run shell commands at agent lifecycle points (tool +// calls, prompt submission, session start/end, etc.) without writing Go or +// compiling a plugin. It is a leaf package depending only on the standard +// library, so it can be composed into the runtime/cli layers without creating +// an import cycle. +// +// This file defines the configuration types and their validation. A hook is a +// single shell command; a matcher binds a group of hooks to an event (and, +// for tool events, a tool-name pattern); a HookSet maps each event type to its +// matchers. The types carry JSON tags so they can live directly in the layered +// config.json. +package hooks + +import ( + "errors" + "fmt" + "io" + "strings" +) + +// DefaultTimeoutSeconds is the per-hook execution timeout when HookConfig.Timeout +// is nil (FR-11). A slow or hung hook is killed after this many seconds and +// treated as a failure (fail-open for blocking hooks). +const DefaultTimeoutSeconds = 60 + +// CommandType is the only hook type supported in v1: a shell command. Other +// types (embedded script engines, WASM) are explicitly out of scope. +const CommandType = "command" + +// HookConfig is a single hook command. +type HookConfig struct { + Type string `json:"type"` // v1: fixed "command" + Command string `json:"command"` // handed to the system shell + Timeout *int `json:"timeout,omitempty"` // seconds; nil = DefaultTimeoutSeconds +} + +// HookMatcherConfig binds a group of hooks to a matcher. An empty (or "*") +// matcher applies to every trigger of the event; otherwise it is matched +// against the tool name (see matcher.go). +type HookMatcherConfig struct { + Matcher string `json:"matcher,omitempty"` + Hooks []HookConfig `json:"hooks"` +} + +// HookSet maps an event type (e.g. "PreToolUse") to its matcher list. It is +// the shape stored in a ConfigLayer and in the resolved Config. +type HookSet map[string][]HookMatcherConfig + +// TimeoutSeconds returns the effective timeout for the hook: its own Timeout +// when set to a positive value, otherwise DefaultTimeoutSeconds. A non-positive +// override is ignored so a misconfigured 0/negative value cannot disable the +// timeout guard. +func (h HookConfig) TimeoutSeconds() int { + if h.Timeout != nil && *h.Timeout > 0 { + return *h.Timeout + } + return DefaultTimeoutSeconds +} + +// Validate reports whether the hook is well-formed: the type must be "command" +// (empty is accepted and treated as "command" for convenience) and the command +// must be non-empty. An invalid hook is rejected at load time and skipped with +// a warning rather than executed. +func (h HookConfig) Validate() error { + if h.Type != "" && h.Type != CommandType { + return errors.New("hook type must be \"command\"") + } + if strings.TrimSpace(h.Command) == "" { + return errors.New("hook command must not be empty") + } + return nil +} + +// warnf writes a formatted warning to w when w is non-nil. Hook failures and +// misconfigurations are surfaced this way (mirroring plugin.EventNotifier's +// warnLog) so a bad hook never interrupts the agent. +func warnf(w io.Writer, format string, args ...any) { + if w == nil { + return + } + fmt.Fprintf(w, format, args...) +} diff --git a/pigo/internal/hooks/config_test.go b/pigo/internal/hooks/config_test.go new file mode 100644 index 0000000..bc2d3c0 --- /dev/null +++ b/pigo/internal/hooks/config_test.go @@ -0,0 +1,47 @@ +package hooks + +import "testing" + +func ptr[T any](v T) *T { return &v } + +func TestHookConfigValidate(t *testing.T) { + tests := []struct { + name string + h HookConfig + wantErr bool + }{ + {"valid command", HookConfig{Type: "command", Command: "echo hi"}, false}, + {"empty type defaults ok", HookConfig{Command: "echo hi"}, false}, + {"empty command", HookConfig{Type: "command", Command: ""}, true}, + {"whitespace command", HookConfig{Type: "command", Command: " "}, true}, + {"wrong type", HookConfig{Type: "wasm", Command: "echo hi"}, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.h.Validate() + if (err != nil) != tt.wantErr { + t.Fatalf("Validate() err = %v, wantErr = %v", err, tt.wantErr) + } + }) + } +} + +func TestHookConfigTimeoutSeconds(t *testing.T) { + tests := []struct { + name string + h HookConfig + want int + }{ + {"nil timeout uses default", HookConfig{}, DefaultTimeoutSeconds}, + {"positive override", HookConfig{Timeout: ptr(10)}, 10}, + {"zero override ignored", HookConfig{Timeout: ptr(0)}, DefaultTimeoutSeconds}, + {"negative override ignored", HookConfig{Timeout: ptr(-5)}, DefaultTimeoutSeconds}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.h.TimeoutSeconds(); got != tt.want { + t.Fatalf("TimeoutSeconds() = %d, want %d", got, tt.want) + } + }) + } +} diff --git a/pigo/internal/hooks/dispatch.go b/pigo/internal/hooks/dispatch.go new file mode 100644 index 0000000..86f517b --- /dev/null +++ b/pigo/internal/hooks/dispatch.go @@ -0,0 +1,91 @@ +// This file implements the Dispatcher: for one event it matches the configured +// hooks, runs them in order via the Runner, and merges their outputs into a +// single HookDecision. It owns the fail-open policy (a hook that fails to run +// is warned about and skipped, never blocking the agent) and the PreToolUse +// short-circuit (the first block stops further hooks so a blocked call is not +// also rewritten). +package hooks + +import ( + "context" + "io" + "strings" +) + +// EventPreToolUse is the one event whose first block short-circuits the rest of +// the chain. Declared here (rather than importing a shared constants file) so +// the hooks package stays a leaf; the string must match what callers dispatch. +const EventPreToolUse = "PreToolUse" + +// Dispatcher runs the hooks configured for an event and merges their results. +// A nil *Dispatcher is a valid no-op (Dispatch returns an empty decision), so +// callers can hold a possibly-nil dispatcher without guarding every call. +type Dispatcher struct { + set HookSet + runner *Runner + warnLog io.Writer +} + +// NewDispatcher builds a dispatcher over the given hook set. It returns nil when +// the set is empty, so the common no-hooks case costs nothing and callers can +// treat nil as "hooks disabled" (FR-18). projectDir is where hook commands run; +// warnLog receives isolation warnings (may be nil). +func NewDispatcher(set HookSet, projectDir string, warnLog io.Writer) *Dispatcher { + if len(set) == 0 { + return nil + } + return &Dispatcher{ + set: set, + runner: &Runner{ProjectDir: projectDir, WarnLog: warnLog}, + warnLog: warnLog, + } +} + +// Dispatch runs every hook matching (eventType, toolName) in order and returns +// the merged decision. On a nil dispatcher or no matched hooks it returns the +// zero HookDecision. A hook that fails to run is warned about and skipped +// (fail-open, FR-15). For PreToolUse the first block stops the chain so a +// blocked call is not subsequently rewritten. +func (d *Dispatcher) Dispatch(ctx context.Context, eventType, toolName string, input HookInput) HookDecision { + var dec HookDecision + if d == nil { + return dec + } + matched := d.set.MatchHooks(eventType, toolName, d.warnLog) + for _, h := range matched { + out, err := d.runner.Run(ctx, h, input) + if err != nil { + warnf(d.warnLog, "pigo: hooks: %s: %v\n", eventType, err) + continue // fail-open + } + if out.blocks() { + dec.Block = true + dec.Reason = joinNonEmpty(dec.Reason, out.Reason, "\n") + } + if out.AdditionalContext != "" { + dec.AdditionalContext = joinNonEmpty(dec.AdditionalContext, out.AdditionalContext, "\n") + } + if len(out.UpdatedInput) > 0 { + dec.UpdatedInput = out.UpdatedInput // last writer wins (§5.4) + } + if dec.Block && eventType == EventPreToolUse { + break // blocked tool call is not also rewritten + } + } + return dec +} + +// joinNonEmpty joins a and b with sep, dropping empty operands so the result +// never has a leading or dangling separator. +func joinNonEmpty(a, b, sep string) string { + a = strings.TrimSpace(a) + b = strings.TrimSpace(b) + switch { + case a == "": + return b + case b == "": + return a + default: + return a + sep + b + } +} diff --git a/pigo/internal/hooks/dispatch_test.go b/pigo/internal/hooks/dispatch_test.go new file mode 100644 index 0000000..db4c958 --- /dev/null +++ b/pigo/internal/hooks/dispatch_test.go @@ -0,0 +1,112 @@ +package hooks + +import ( + "bytes" + "context" + "runtime" + "strings" + "testing" +) + +func TestNewDispatcherNilOnEmpty(t *testing.T) { + if d := NewDispatcher(nil, "/tmp", nil); d != nil { + t.Fatal("expected nil dispatcher for empty set") + } + if d := NewDispatcher(HookSet{}, "/tmp", nil); d != nil { + t.Fatal("expected nil dispatcher for empty set") + } +} + +func TestNilDispatcherDispatchIsNoOp(t *testing.T) { + var d *Dispatcher + dec := d.Dispatch(context.Background(), "PreToolUse", "bash", HookInput{}) + if dec.Block || dec.AdditionalContext != "" || dec.UpdatedInput != nil { + t.Fatalf("expected empty decision, got %+v", dec) + } +} + +func TestDispatchMergesContext(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + set := HookSet{ + "PostToolUse": { + {Matcher: "*", Hooks: []HookConfig{ + {Command: `echo '{"additionalContext":"first"}'`}, + {Command: `echo '{"additionalContext":"second"}'`}, + }}, + }, + } + d := NewDispatcher(set, t.TempDir(), nil) + dec := d.Dispatch(context.Background(), "PostToolUse", "write", HookInput{}) + if dec.AdditionalContext != "first\nsecond" { + t.Fatalf("expected merged context, got %q", dec.AdditionalContext) + } +} + +func TestDispatchPreToolUseFirstBlockShortCircuits(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + set := HookSet{ + "PreToolUse": { + {Matcher: "*", Hooks: []HookConfig{ + {Command: `echo "stop" >&2; exit 2`}, + {Command: `echo '{"updatedInput":{"changed":true}}'`}, + }}, + }, + } + d := NewDispatcher(set, t.TempDir(), nil) + dec := d.Dispatch(context.Background(), "PreToolUse", "bash", HookInput{}) + if !dec.Block { + t.Fatal("expected block") + } + if dec.UpdatedInput != nil { + t.Fatalf("expected short-circuit before rewrite, got %s", dec.UpdatedInput) + } +} + +func TestDispatchFailOpen(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + set := HookSet{ + "PreToolUse": { + {Matcher: "*", Hooks: []HookConfig{ + {Command: `exit 1`}, + {Command: `echo '{"additionalContext":"survived"}'`}, + }}, + }, + } + var warn bytes.Buffer + d := NewDispatcher(set, t.TempDir(), &warn) + dec := d.Dispatch(context.Background(), "PreToolUse", "bash", HookInput{}) + if dec.Block { + t.Fatal("failed hook must not block (fail-open)") + } + if dec.AdditionalContext != "survived" { + t.Fatalf("expected later hook to still run, got %q", dec.AdditionalContext) + } + if !strings.Contains(warn.String(), "PreToolUse") { + t.Fatalf("expected failure warning, got %q", warn.String()) + } +} + +func TestDispatchUpdatedInputLastWins(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + set := HookSet{ + "PreToolUse": { + {Matcher: "*", Hooks: []HookConfig{ + {Command: `echo '{"updatedInput":{"n":1}}'`}, + {Command: `echo '{"updatedInput":{"n":2}}'`}, + }}, + }, + } + d := NewDispatcher(set, t.TempDir(), nil) + dec := d.Dispatch(context.Background(), "PreToolUse", "bash", HookInput{}) + if string(dec.UpdatedInput) != `{"n":2}` { + t.Fatalf("expected last writer wins, got %s", dec.UpdatedInput) + } +} diff --git a/pigo/internal/hooks/matcher.go b/pigo/internal/hooks/matcher.go new file mode 100644 index 0000000..b38ef8e --- /dev/null +++ b/pigo/internal/hooks/matcher.go @@ -0,0 +1,85 @@ +// This file implements hook matching: given an event type and (for tool +// events) a tool name, MatchHooks returns the flat, ordered list of hooks that +// should run. Matcher semantics follow Claude Code to keep the learning curve +// low: empty or "*" matches all, an exact tool name matches only that tool, a +// "|"-separated list matches any listed tool, and anything else is compiled as +// a Go regexp against the tool name. +package hooks + +import ( + "io" + "regexp" + "strings" +) + +// MatchHooks returns every valid hook under eventType whose matcher matches +// toolName, preserving config order (layer order + declaration order within a +// layer). For events that do not carry a tool name (toolName == ""), matchers +// are ignored and all hooks under the event fire. Invalid hooks (see +// HookConfig.Validate) are skipped; a matcher whose regexp fails to compile is +// skipped with a warning on warnLog (when non-nil). +func (s HookSet) MatchHooks(eventType, toolName string, warnLog io.Writer) []HookConfig { + matchers := s[eventType] + if len(matchers) == 0 { + return nil + } + var out []HookConfig + for _, m := range matchers { + if !matcherApplies(m.Matcher, toolName, warnLog) { + continue + } + for _, h := range m.Hooks { + if err := h.Validate(); err != nil { + warnf(warnLog, "pigo: hooks: skipping invalid hook: %v\n", err) + continue + } + out = append(out, h) + } + } + return out +} + +// matcherApplies reports whether a matcher pattern matches the given tool name. +// An empty tool name (event without a tool) always matches, so tool-agnostic +// events fire every hook regardless of matcher. +func matcherApplies(pattern, toolName string, warnLog io.Writer) bool { + if toolName == "" { + return true + } + pattern = strings.TrimSpace(pattern) + if pattern == "" || pattern == "*" { + return true + } + // "|"-separated multi-value: any exact tool name matches. This also covers + // the single-exact-name case (no "|"). + parts := strings.Split(pattern, "|") + exactCandidate := true + for _, p := range parts { + p = strings.TrimSpace(p) + if p == toolName { + return true + } + if !isPlainToolName(p) { + exactCandidate = false + } + } + // If every alternative was a plain tool name, this was an exact/multi-value + // matcher that simply did not match — do not fall through to regexp. + if exactCandidate { + return false + } + re, err := regexp.Compile(pattern) + if err != nil { + warnf(warnLog, "pigo: hooks: skipping matcher with invalid regexp %q: %v\n", pattern, err) + return false + } + return re.MatchString(toolName) +} + +// isPlainToolName reports whether s looks like a literal tool name rather than +// a regexp — i.e. it contains no regexp metacharacters. Used to decide whether +// a "|"-split alternative should be treated as an exact match or as part of a +// regexp alternation. +func isPlainToolName(s string) bool { + return !strings.ContainsAny(s, ".*+?()[]{}^$\\") +} diff --git a/pigo/internal/hooks/matcher_test.go b/pigo/internal/hooks/matcher_test.go new file mode 100644 index 0000000..be0ab78 --- /dev/null +++ b/pigo/internal/hooks/matcher_test.go @@ -0,0 +1,111 @@ +package hooks + +import ( + "bytes" + "strings" + "testing" +) + +func TestMatchHooks(t *testing.T) { + set := HookSet{ + "PreToolUse": { + {Matcher: "", Hooks: []HookConfig{{Command: "all"}}}, + {Matcher: "*", Hooks: []HookConfig{{Command: "star"}}}, + {Matcher: "bash", Hooks: []HookConfig{{Command: "bash-only"}}}, + {Matcher: "write|edit", Hooks: []HookConfig{{Command: "write-or-edit"}}}, + {Matcher: "Edit.*", Hooks: []HookConfig{{Command: "regex-edit"}}}, + }, + "SessionStart": { + {Matcher: "ignored", Hooks: []HookConfig{{Command: "session"}}}, + }, + } + + commands := func(hs []HookConfig) []string { + out := make([]string, len(hs)) + for i, h := range hs { + out[i] = h.Command + } + return out + } + + t.Run("bash matches empty, star, exact", func(t *testing.T) { + got := commands(set.MatchHooks("PreToolUse", "bash", nil)) + want := []string{"all", "star", "bash-only"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Fatalf("got %v, want %v", got, want) + } + }) + + t.Run("write matches multi-value", func(t *testing.T) { + got := commands(set.MatchHooks("PreToolUse", "write", nil)) + if !contains(got, "write-or-edit") { + t.Fatalf("expected write-or-edit in %v", got) + } + }) + + t.Run("EditFile matches regex not exact", func(t *testing.T) { + got := commands(set.MatchHooks("PreToolUse", "EditFile", nil)) + if !contains(got, "regex-edit") { + t.Fatalf("expected regex-edit in %v", got) + } + if contains(got, "bash-only") { + t.Fatalf("did not expect bash-only in %v", got) + } + }) + + t.Run("no-tool-name event ignores matcher", func(t *testing.T) { + got := commands(set.MatchHooks("SessionStart", "", nil)) + want := []string{"session"} + if strings.Join(got, ",") != strings.Join(want, ",") { + t.Fatalf("got %v, want %v", got, want) + } + }) + + t.Run("unknown event returns nil", func(t *testing.T) { + if got := set.MatchHooks("Nope", "bash", nil); got != nil { + t.Fatalf("expected nil, got %v", got) + } + }) +} + +func TestMatchHooksInvalidRegexSkipped(t *testing.T) { + set := HookSet{ + "PreToolUse": { + {Matcher: "(unterminated", Hooks: []HookConfig{{Command: "bad"}}}, + {Matcher: "bash", Hooks: []HookConfig{{Command: "good"}}}, + }, + } + var warn bytes.Buffer + got := set.MatchHooks("PreToolUse", "bash", &warn) + if len(got) != 1 || got[0].Command != "good" { + t.Fatalf("expected only good hook, got %v", got) + } + if !strings.Contains(warn.String(), "invalid regexp") { + t.Fatalf("expected regexp warning, got %q", warn.String()) + } +} + +func TestMatchHooksSkipsInvalidHook(t *testing.T) { + set := HookSet{ + "PreToolUse": { + {Matcher: "*", Hooks: []HookConfig{{Command: ""}, {Command: "ok"}}}, + }, + } + var warn bytes.Buffer + got := set.MatchHooks("PreToolUse", "bash", &warn) + if len(got) != 1 || got[0].Command != "ok" { + t.Fatalf("expected only ok hook, got %v", got) + } + if !strings.Contains(warn.String(), "invalid hook") { + t.Fatalf("expected invalid-hook warning, got %q", warn.String()) + } +} + +func contains(xs []string, s string) bool { + for _, x := range xs { + if x == s { + return true + } + } + return false +} diff --git a/pigo/internal/hooks/notifier.go b/pigo/internal/hooks/notifier.go new file mode 100644 index 0000000..b16e445 --- /dev/null +++ b/pigo/internal/hooks/notifier.go @@ -0,0 +1,118 @@ +// This file bridges the agent's event stream to observer-only hooks (US-011/012/ +// 013, FR-1). Where PreToolUse/UserPromptSubmit/Stop are decision hooks wired at +// dedicated seams, SessionEnd/PreCompact/Notification only observe: the agent +// emits lifecycle events, HookNotifier maps each to a HookInput and fires the +// matching hooks via the Dispatcher, discarding the decision (an observer hook +// cannot block). +// +// It mirrors plugin.EventNotifier: created once per run, its Handle method is +// wired as an event-stream OnEvent callback and coexists with the plugin +// notifier (RunConfig.OnEvent already chains multiple observers). A nil +// *Dispatcher makes NewHookNotifier return nil, and every method is a no-op on a +// nil receiver, so callers can wire it unconditionally. +// +// Session id and project dir are fixed for a run, so they are captured at +// construction rather than read from each event (AgentEndEvent/CompactionEvent +// do not carry them). +package hooks + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// HookNotifier forwards agent lifecycle events to observer-only hooks. Handle +// maps AgentEndEvent→SessionEnd and CompactionEvent→PreCompact; Notify emits a +// Notification event for out-of-band prompts (e.g. a trust confirmation). The +// merged decision is intentionally discarded — these events have no in-flight +// action to veto. +type HookNotifier struct { + d *Dispatcher + sessionID string + projectDir string +} + +// NewHookNotifier returns a notifier over d, or nil when d is nil (no hooks) so +// the caller can skip the OnEvent wiring entirely. sessionID and projectDir +// populate every emitted HookInput. +func NewHookNotifier(d *Dispatcher, sessionID, projectDir string) *HookNotifier { + if d == nil { + return nil + } + return &HookNotifier{d: d, sessionID: sessionID, projectDir: projectDir} +} + +// Handle maps an observed event to its observer hook and fires it. Events with +// no observer mapping (turn/message/tool events) are ignored. It is a no-op on a +// nil notifier, so it can be chained onto OnEvent unconditionally. +func (n *HookNotifier) Handle(ev agentcore.AgentEvent) { + if n == nil { + return + } + switch e := ev.(type) { + case agentcore.AgentEndEvent: + n.dispatch("SessionEnd", HookInput{ + EventType: "SessionEnd", + SessionID: n.sessionID, + ProjectDir: n.projectDir, + StopReason: sessionEndReason(e.Messages), + }) + case agentcore.CompactionEvent: + n.dispatch("PreCompact", HookInput{ + EventType: "PreCompact", + SessionID: n.sessionID, + ProjectDir: n.projectDir, + Trigger: compactionTrigger(e.Reason), + }) + } +} + +// Notify fires the Notification event with a human-readable message. It is used +// for out-of-band prompts that are not part of the event stream, such as a trust +// confirmation for an untrusted-directory bash/write. A no-op on a nil notifier +// or an empty message. +func (n *HookNotifier) Notify(message string) { + if n == nil || message == "" { + return + } + n.dispatch("Notification", HookInput{ + EventType: "Notification", + SessionID: n.sessionID, + ProjectDir: n.projectDir, + Message: message, + }) +} + +// dispatch fires the hooks for an observer event and discards the decision. +func (n *HookNotifier) dispatch(event string, input HookInput) { + n.d.Dispatch(context.Background(), event, "", input) +} + +// sessionEndReason derives the SessionEnd reason from the run's terminal +// assistant message: "error"/"aborted" pass through, everything else (end_turn/ +// tool_use/length or no assistant message) is a "natural" end. +func sessionEndReason(msgs []agentcore.AgentMessage) string { + for i := len(msgs) - 1; i >= 0; i-- { + if am, ok := msgs[i].(agentcore.AssistantMessage); ok { + switch am.StopReason { + case agentcore.StopReasonError: + return "error" + case agentcore.StopReasonAborted: + return "aborted" + default: + return "natural" + } + } + } + return "natural" +} + +// compactionTrigger maps a CompactionEvent.Reason to the PreCompact trigger: +// "manual" stays manual; "threshold"/"overflow" (and anything else) are "auto". +func compactionTrigger(reason string) string { + if reason == "manual" { + return "manual" + } + return "auto" +} diff --git a/pigo/internal/hooks/notifier_test.go b/pigo/internal/hooks/notifier_test.go new file mode 100644 index 0000000..55dd93c --- /dev/null +++ b/pigo/internal/hooks/notifier_test.go @@ -0,0 +1,122 @@ +package hooks + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// captureNotifier builds a HookNotifier whose hook for `event` appends its stdin +// payload to a file, returning the notifier and the capture-file path so a test +// can assert on the payload the hook received. +func captureNotifier(t *testing.T, event string) (*HookNotifier, string) { + t.Helper() + dir := t.TempDir() + out := filepath.Join(dir, "capture.json") + set := HookSet{ + event: {{Matcher: "*", Hooks: []HookConfig{{Command: "cat >> " + out}}}}, + } + d := NewDispatcher(set, dir, nil) + if d == nil { + t.Fatal("expected non-nil dispatcher") + } + return NewHookNotifier(d, "sess-1", dir), out +} + +func readCapture(t *testing.T, path string) string { + t.Helper() + b, err := os.ReadFile(path) + if err != nil { + t.Fatalf("capture not written: %v", err) + } + return string(b) +} + +// TestSessionEndNaturalReason: an end_turn terminal message yields reason +// "natural"; an aborted terminal message yields "aborted". +func TestSessionEndReason(t *testing.T) { + cases := []struct { + name string + stop string + expect string + }{ + {"natural", agentcore.StopReasonEndTurn, `"stop_reason":"natural"`}, + {"aborted", agentcore.StopReasonAborted, `"stop_reason":"aborted"`}, + {"error", agentcore.StopReasonError, `"stop_reason":"error"`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + n, out := captureNotifier(t, "SessionEnd") + n.Handle(agentcore.AgentEndEvent{Messages: []agentcore.AgentMessage{ + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: tc.stop}, + }}) + if got := readCapture(t, out); !strings.Contains(got, tc.expect) { + t.Fatalf("SessionEnd reason: want %q in %q", tc.expect, got) + } + }) + } +} + +// TestPreCompactTrigger: a manual CompactionEvent maps to trigger "manual"; +// threshold/overflow map to "auto". +func TestPreCompactTrigger(t *testing.T) { + cases := []struct { + reason string + expect string + }{ + {"manual", `"trigger":"manual"`}, + {"threshold", `"trigger":"auto"`}, + {"overflow", `"trigger":"auto"`}, + } + for _, tc := range cases { + t.Run(tc.reason, func(t *testing.T) { + n, out := captureNotifier(t, "PreCompact") + n.Handle(agentcore.CompactionEvent{Reason: tc.reason}) + if got := readCapture(t, out); !strings.Contains(got, tc.expect) { + t.Fatalf("PreCompact trigger: want %q in %q", tc.expect, got) + } + }) + } +} + +// TestNotification: Notify fires the Notification event carrying the message. +func TestNotification(t *testing.T) { + n, out := captureNotifier(t, "Notification") + n.Notify("approve bash in untrusted dir?") + got := readCapture(t, out) + if !strings.Contains(got, `"event_type":"Notification"`) || !strings.Contains(got, "approve bash in untrusted dir?") { + t.Fatalf("Notification payload missing message, got %q", got) + } +} + +// TestNotifierNilSafe: a nil dispatcher yields a nil notifier whose methods are +// safe no-ops; an empty Notify message is dropped. +func TestNotifierNilSafe(t *testing.T) { + if n := NewHookNotifier(nil, "s", "d"); n != nil { + t.Fatal("nil dispatcher must yield a nil notifier") + } + var n *HookNotifier + n.Handle(agentcore.AgentEndEvent{}) + n.Notify("x") // must not panic + + // A live notifier with an empty message writes nothing. + live, out := captureNotifier(t, "Notification") + live.Notify("") + if _, err := os.Stat(out); !os.IsNotExist(err) { + t.Fatal("empty Notify message must not fire the hook") + } +} + +// TestNotifierIgnoresUnmappedEvents: turn/message/tool events have no observer +// mapping and must not fire any hook. +func TestNotifierIgnoresUnmappedEvents(t *testing.T) { + n, out := captureNotifier(t, "SessionEnd") + n.Handle(agentcore.TurnStartEvent{}) + n.Handle(agentcore.ToolExecutionStartEvent{ToolName: "bash"}) + if _, err := os.Stat(out); !os.IsNotExist(err) { + t.Fatal("unmapped events must not fire the SessionEnd hook") + } +} diff --git a/pigo/internal/hooks/protocol.go b/pigo/internal/hooks/protocol.go new file mode 100644 index 0000000..e27c665 --- /dev/null +++ b/pigo/internal/hooks/protocol.go @@ -0,0 +1,82 @@ +// This file defines the process contract between pigo and a user hook command: +// what pigo writes to the hook's stdin (HookInput), what pigo reads back from +// its stdout (HookOutput), and the internal merged decision the dispatcher +// produces (HookDecision). It also defines the exit-code semantics parsing. +// +// The wire contract follows Claude Code: exit 0 = allow, exit 2 = block (with +// stderr as the reason), any other non-zero = execution failure. On exit 0 a +// well-formed JSON stdout is parsed as HookOutput; a non-JSON stdout is a no-op. +package hooks + +import ( + "encoding/json" + "strings" +) + +// HookInput is the JSON payload pigo writes to a hook command's stdin. It +// carries only observable, non-secret fields (FR-17) — never API keys or +// credentials. Per-event fields are omitempty so a payload only contains the +// fields relevant to its event type. +type HookInput struct { + EventType string `json:"event_type"` + SessionID string `json:"session_id,omitempty"` + ProjectDir string `json:"project_dir,omitempty"` + ToolName string `json:"tool_name,omitempty"` // Pre/PostToolUse + ToolInput json.RawMessage `json:"tool_input,omitempty"` // Pre/PostToolUse + ToolResponse json.RawMessage `json:"tool_response,omitempty"` // PostToolUse + Prompt string `json:"prompt,omitempty"` // UserPromptSubmit + StopReason string `json:"stop_reason,omitempty"` // Stop/SessionEnd + Source string `json:"source,omitempty"` // SessionStart (startup/resume) + Trigger string `json:"trigger,omitempty"` // PreCompact (manual/auto) + Message string `json:"message,omitempty"` // Notification +} + +// HookOutput is the optional JSON a hook may print to stdout to influence the +// agent. A non-JSON stdout on exit 0 is treated as an empty HookOutput (no +// operation). Decision "block" is equivalent to exiting with code 2. +type HookOutput struct { + Decision string `json:"decision,omitempty"` // "block" | "approve" | "" + Reason string `json:"reason,omitempty"` + AdditionalContext string `json:"additionalContext,omitempty"` + Continue *bool `json:"continue,omitempty"` + UpdatedInput json.RawMessage `json:"updatedInput,omitempty"` // PreToolUse: rewrite tool args +} + +// blocks reports whether this output requests a block. A decision of "block" +// blocks; an explicit continue=false also blocks. "approve" and the empty +// decision allow. +func (o HookOutput) blocks() bool { + if strings.EqualFold(o.Decision, "block") { + return true + } + if o.Continue != nil && !*o.Continue { + return true + } + return false +} + +// HookDecision is the dispatcher's merged result after running all matched +// hooks for one event. Block is set if any hook blocked; Reason accumulates the +// blocking reasons; AdditionalContext accumulates injected context in order; +// UpdatedInput holds the last-provided rewrite (last writer wins, §5.4). +type HookDecision struct { + Block bool + Reason string + AdditionalContext string + UpdatedInput json.RawMessage +} + +// parseHookOutput parses a hook's stdout into a HookOutput. On exit 0 a +// non-JSON body is a no-op (returns the zero value, ok=false). Empty/whitespace +// stdout is also a no-op. A valid JSON object is parsed and ok is true. +func parseHookOutput(stdout []byte) (HookOutput, bool) { + trimmed := strings.TrimSpace(string(stdout)) + if trimmed == "" { + return HookOutput{}, false + } + var out HookOutput + if err := json.Unmarshal([]byte(trimmed), &out); err != nil { + return HookOutput{}, false + } + return out, true +} diff --git a/pigo/internal/hooks/protocol_test.go b/pigo/internal/hooks/protocol_test.go new file mode 100644 index 0000000..b1dcc92 --- /dev/null +++ b/pigo/internal/hooks/protocol_test.go @@ -0,0 +1,80 @@ +package hooks + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestParseHookOutput(t *testing.T) { + t.Run("empty is no-op", func(t *testing.T) { + if _, ok := parseHookOutput([]byte(" ")); ok { + t.Fatal("expected ok=false for empty") + } + }) + t.Run("non-json is no-op", func(t *testing.T) { + if _, ok := parseHookOutput([]byte("just some text")); ok { + t.Fatal("expected ok=false for non-json") + } + }) + t.Run("valid decision block", func(t *testing.T) { + out, ok := parseHookOutput([]byte(`{"decision":"block","reason":"nope"}`)) + if !ok || !out.blocks() || out.Reason != "nope" { + t.Fatalf("unexpected: ok=%v out=%+v", ok, out) + } + }) + t.Run("additionalContext", func(t *testing.T) { + out, ok := parseHookOutput([]byte(`{"additionalContext":"extra"}`)) + if !ok || out.AdditionalContext != "extra" { + t.Fatalf("unexpected: ok=%v out=%+v", ok, out) + } + }) + t.Run("updatedInput preserved as raw", func(t *testing.T) { + out, ok := parseHookOutput([]byte(`{"updatedInput":{"a":1}}`)) + if !ok || string(out.UpdatedInput) != `{"a":1}` { + t.Fatalf("unexpected: ok=%v raw=%s", ok, out.UpdatedInput) + } + }) +} + +func TestHookInputNoSecretFields(t *testing.T) { + // A fully-populated payload must never carry credential-like keys: the + // struct is a whitelist, so marshaling can only emit its declared fields. + in := HookInput{ + EventType: "PreToolUse", SessionID: "s", ProjectDir: "/p", ToolName: "bash", + ToolInput: json.RawMessage(`{"cmd":"ls"}`), Prompt: "hi", Message: "m", + } + data, err := json.Marshal(in) + if err != nil { + t.Fatalf("marshal: %v", err) + } + forbidden := []string{"api_key", "apikey", "token", "credential", "password", "secret", "authorization"} + lower := strings.ToLower(string(data)) + for _, k := range forbidden { + if strings.Contains(lower, k) { + t.Fatalf("payload contains forbidden key %q: %s", k, data) + } + } +} + +func TestHookOutputBlocks(t *testing.T) { + tests := []struct { + name string + out HookOutput + want bool + }{ + {"decision block", HookOutput{Decision: "block"}, true}, + {"decision BLOCK case-insensitive", HookOutput{Decision: "BLOCK"}, true}, + {"decision approve", HookOutput{Decision: "approve"}, false}, + {"empty decision", HookOutput{}, false}, + {"continue false blocks", HookOutput{Continue: ptr(false)}, true}, + {"continue true allows", HookOutput{Continue: ptr(true)}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.out.blocks(); got != tt.want { + t.Fatalf("blocks() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/pigo/internal/hooks/runner.go b/pigo/internal/hooks/runner.go new file mode 100644 index 0000000..89ccaff --- /dev/null +++ b/pigo/internal/hooks/runner.go @@ -0,0 +1,162 @@ +// This file implements Runner.Run: executing a single hook command via the +// system shell with the event payload on stdin, a per-hook timeout, bounded +// output capture, and exit-code classification. It is the one place that forks +// a process, so all isolation guarantees (timeout kill, output cap, error +// containment) live here. +package hooks + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "os/exec" + "strings" + "time" +) + +// MaxOutputBytes caps the bytes captured from a hook's stdout and stderr each +// (FR-13). Output beyond this is dropped and the truncation is flagged; the +// captured prefix is still parsed so a hook that prints a small JSON decision +// followed by noise still works. +const MaxOutputBytes = 1 << 20 // 1 MB + +// blockExitCode is the exit code that signals a block (Claude Code semantics): +// the command exited 2, stderr carries the reason. +const blockExitCode = 2 + +// Runner executes hook commands. Shell defaults to "sh" with a "-c" flag; it is +// a field so tests can substitute a shell and future platforms can override it. +// ProjectDir is the working directory hook commands run in (the project root). +// ExtraEnv is appended to the process environment (PIGO_* variables). +type Runner struct { + Shell string + ProjectDir string + WarnLog io.Writer +} + +// Run executes one hook, writing input as a single-line JSON document to the +// command's stdin. It returns the parsed HookOutput and a non-nil error only +// for an execution *failure* (could not start, timed out, or exited non-zero +// and non-2). A clean exit 0 or a block (exit 2) both return err == nil; the +// caller distinguishes a block via HookOutput.blocks(). The PIGO_* environment +// variables are injected on top of the current process environment. +func (r *Runner) Run(ctx context.Context, h HookConfig, input HookInput) (HookOutput, error) { + payload, err := json.Marshal(input) + if err != nil { + return HookOutput{}, fmt.Errorf("marshal hook input: %w", err) + } + + timeout := time.Duration(h.TimeoutSeconds()) * time.Second + runCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + shell := r.Shell + if shell == "" { + shell = "sh" + } + cmd := exec.CommandContext(runCtx, shell, "-c", h.Command) + cmd.Dir = r.ProjectDir + cmd.Env = r.env(input) + cmd.Stdin = bytes.NewReader(payload) + // On timeout, CommandContext kills the shell, but a grandchild it spawned + // (e.g. `sh -c "sleep 5"` forking sleep) can inherit the stdout/stderr pipes + // and keep them open, blocking cmd.Run until that grandchild exits on its own + // — so Run would return only after the full command duration, not the timeout. + // WaitDelay bounds that wait: after the process is killed, Go force-closes the + // I/O pipes so Run returns promptly. + cmd.WaitDelay = time.Second + + var stdout, stderr cappedBuffer + stdout.limit = MaxOutputBytes + stderr.limit = MaxOutputBytes + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + runErr := cmd.Run() + + if stdout.truncated() || stderr.truncated() { + warnf(r.WarnLog, "pigo: hooks: output from command %q exceeded %d bytes and was truncated\n", h.Command, MaxOutputBytes) + } + + // Timeout: the context deadline fired and the process was killed. + if runCtx.Err() == context.DeadlineExceeded { + return HookOutput{}, fmt.Errorf("hook timed out after %s", timeout) + } + + exitCode := 0 + if runErr != nil { + var ee *exec.ExitError + if errors.As(runErr, &ee) { + exitCode = ee.ExitCode() + } else { + // Could not start (ENOENT etc.) or was killed. + return HookOutput{}, fmt.Errorf("hook failed to run: %w", runErr) + } + } + + switch exitCode { + case 0: + out, _ := parseHookOutput(stdout.Bytes()) + return out, nil + case blockExitCode: + // Block: prefer a JSON decision if present, else synthesize one from + // stderr as the reason. + if out, ok := parseHookOutput(stdout.Bytes()); ok { + if out.Reason == "" { + out.Reason = strings.TrimSpace(string(stderr.Bytes())) + } + out.Decision = "block" + return out, nil + } + return HookOutput{Decision: "block", Reason: strings.TrimSpace(string(stderr.Bytes()))}, nil + default: + return HookOutput{}, fmt.Errorf("hook exited with code %d: %s", exitCode, strings.TrimSpace(string(stderr.Bytes()))) + } +} + +// env builds the command environment: the current process environment plus the +// PIGO_* variables derived from the input. +func (r *Runner) env(input HookInput) []string { + env := append([]string(nil), os.Environ()...) + env = append(env, + "PIGO_SESSION_ID="+input.SessionID, + "PIGO_PROJECT_DIR="+r.ProjectDir, + "PIGO_EVENT_TYPE="+input.EventType, + ) + return env +} + +// cappedBuffer is an io.Writer that stores at most limit bytes and counts how +// many it dropped, so hook output cannot exhaust memory (FR-13). +type cappedBuffer struct { + buf bytes.Buffer + limit int + dropped int +} + +func (c *cappedBuffer) Write(p []byte) (int, error) { + if c.limit <= 0 { + return c.buf.Write(p) + } + room := c.limit - c.buf.Len() + if room <= 0 { + c.dropped += len(p) + return len(p), nil + } + if len(p) > room { + c.buf.Write(p[:room]) + c.dropped += len(p) - room + return len(p), nil + } + return c.buf.Write(p) +} + +// Bytes returns the captured (possibly truncated) output. +func (c *cappedBuffer) Bytes() []byte { return c.buf.Bytes() } + +// truncated reports whether any output was dropped by the cap. +func (c *cappedBuffer) truncated() bool { return c.dropped > 0 } diff --git a/pigo/internal/hooks/runner_test.go b/pigo/internal/hooks/runner_test.go new file mode 100644 index 0000000..04dac2a --- /dev/null +++ b/pigo/internal/hooks/runner_test.go @@ -0,0 +1,130 @@ +package hooks + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" +) + +func TestRunnerRunStdinAndEnv(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + dir := t.TempDir() + outFile := filepath.Join(dir, "captured.json") + envFile := filepath.Join(dir, "env.txt") + + r := &Runner{ProjectDir: dir} + h := HookConfig{Command: "cat > " + outFile + "; printf '%s\\n%s\\n%s\\n' \"$PIGO_SESSION_ID\" \"$PIGO_PROJECT_DIR\" \"$PIGO_EVENT_TYPE\" > " + envFile} + input := HookInput{EventType: "PreToolUse", SessionID: "sess-1", ProjectDir: dir, ToolName: "bash"} + + if _, err := r.Run(context.Background(), h, input); err != nil { + t.Fatalf("Run() error: %v", err) + } + + data, err := os.ReadFile(outFile) + if err != nil { + t.Fatalf("read captured stdin: %v", err) + } + var got HookInput + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("stdin was not valid JSON: %v (%s)", err, data) + } + if got.EventType != "PreToolUse" || got.SessionID != "sess-1" || got.ToolName != "bash" { + t.Fatalf("unexpected decoded stdin: %+v", got) + } + + envData, _ := os.ReadFile(envFile) + lines := strings.Split(strings.TrimSpace(string(envData)), "\n") + if len(lines) != 3 || lines[0] != "sess-1" || lines[1] != dir || lines[2] != "PreToolUse" { + t.Fatalf("unexpected env: %v", lines) + } +} + +func TestRunnerExitCodes(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + dir := t.TempDir() + r := &Runner{ProjectDir: dir} + ctx := context.Background() + + t.Run("exit 0 with json", func(t *testing.T) { + out, err := r.Run(ctx, HookConfig{Command: `echo '{"additionalContext":"hi"}'`}, HookInput{}) + if err != nil { + t.Fatalf("unexpected err: %v", err) + } + if out.AdditionalContext != "hi" { + t.Fatalf("unexpected out: %+v", out) + } + }) + + t.Run("exit 0 non-json is no-op", func(t *testing.T) { + out, err := r.Run(ctx, HookConfig{Command: `echo hello world`}, HookInput{}) + if err != nil || out.blocks() || out.AdditionalContext != "" { + t.Fatalf("expected no-op, got out=%+v err=%v", out, err) + } + }) + + t.Run("exit 2 blocks with stderr reason", func(t *testing.T) { + out, err := r.Run(ctx, HookConfig{Command: `echo "denied" >&2; exit 2`}, HookInput{}) + if err != nil { + t.Fatalf("exit 2 should not be an error, got %v", err) + } + if !out.blocks() || out.Reason != "denied" { + t.Fatalf("unexpected out: %+v", out) + } + }) + + t.Run("exit 1 is failure", func(t *testing.T) { + _, err := r.Run(ctx, HookConfig{Command: `echo boom >&2; exit 1`}, HookInput{}) + if err == nil { + t.Fatal("expected error for exit 1") + } + }) + + t.Run("command not found is failure", func(t *testing.T) { + _, err := r.Run(ctx, HookConfig{Command: `this-command-does-not-exist-pigo`}, HookInput{}) + if err == nil { + t.Fatal("expected error for missing command") + } + }) +} + +func TestRunnerTimeout(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + r := &Runner{ProjectDir: t.TempDir()} + start := time.Now() + _, err := r.Run(context.Background(), HookConfig{Command: "sleep 5", Timeout: ptr(1)}, HookInput{}) + if err == nil || !strings.Contains(err.Error(), "timed out") { + t.Fatalf("expected timeout error, got %v", err) + } + if time.Since(start) > 3*time.Second { + t.Fatalf("timeout took too long: %v", time.Since(start)) + } +} + +func TestRunnerOutputCap(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sh -c not available on windows") + } + var buf cappedBuffer + buf.limit = 10 + n, _ := buf.Write([]byte("0123456789abcdef")) + if n != 16 { + t.Fatalf("Write should report full length, got %d", n) + } + if len(buf.Bytes()) != 10 { + t.Fatalf("expected 10 bytes retained, got %d", len(buf.Bytes())) + } + if !buf.truncated() { + t.Fatal("expected truncated to be true") + } +} diff --git a/pigo/internal/jsonrpc/message.go b/pigo/internal/jsonrpc/message.go new file mode 100644 index 0000000..5579414 --- /dev/null +++ b/pigo/internal/jsonrpc/message.go @@ -0,0 +1,111 @@ +// Package jsonrpc implements a minimal JSON-RPC 2.0 client over a subprocess's +// stdio (US-014/#116). It is the shared transport foundation reused by the MCP +// client (#130/#131), the plugin system (#132/#133) and process-isolated +// sub-agents (#135): each spawns an external executable and speaks line-delimited +// JSON-RPC 2.0 over the child's stdin/stdout. +// +// The wire format follows the spec: every message carries "jsonrpc":"2.0". A +// request has an id and expects a matching response; a notification omits the id +// and expects none. Requests and responses are correlated by id, so concurrent +// requests from different goroutines are safe — each waits only on its own reply. +// +// This file defines the message envelope and (de)serialization; transport.go +// implements the subprocess client. +package jsonrpc + +import ( + "encoding/json" + "fmt" +) + +// Version is the only JSON-RPC protocol version this package speaks. +const Version = "2.0" + +// ID is a JSON-RPC request identifier. The spec allows a string or a number; +// this client only ever generates numeric ids, but ID round-trips whatever a +// peer sends back so response correlation still works against servers that echo +// string ids. +type ID struct { + num int64 + str string + isStr bool +} + +// NumID returns a numeric request id. +func NumID(n int64) ID { return ID{num: n} } + +// String renders the id for use as a map key when correlating responses. +func (id ID) String() string { + if id.isStr { + return "s:" + id.str + } + return fmt.Sprintf("n:%d", id.num) +} + +// MarshalJSON emits the id as its underlying JSON scalar (number or string). +func (id ID) MarshalJSON() ([]byte, error) { + if id.isStr { + return json.Marshal(id.str) + } + return json.Marshal(id.num) +} + +// UnmarshalJSON accepts either a JSON number or string id. +func (id *ID) UnmarshalJSON(data []byte) error { + var n int64 + if err := json.Unmarshal(data, &n); err == nil { + id.num, id.isStr, id.str = n, false, "" + return nil + } + var s string + if err := json.Unmarshal(data, &s); err == nil { + id.str, id.isStr, id.num = s, true, 0 + return nil + } + return fmt.Errorf("jsonrpc: id is neither number nor string: %s", data) +} + +// Request is an outgoing JSON-RPC request or notification. When ID is nil the +// message is a notification (no response expected). +type Request struct { + JSONRPC string `json:"jsonrpc"` + ID *ID `json:"id,omitempty"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` +} + +// Response is an incoming JSON-RPC response. Exactly one of Result / Error is +// set on a well-formed reply. +type Response struct { + JSONRPC string `json:"jsonrpc"` + ID *ID `json:"id,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *Error `json:"error,omitempty"` +} + +// Error is a JSON-RPC error object. +type Error struct { + Code int `json:"code"` + Message string `json:"message"` + Data json.RawMessage `json:"data,omitempty"` +} + +// Error implements the error interface so a peer error can be returned directly. +func (e *Error) Error() string { + return fmt.Sprintf("jsonrpc: server error %d: %s", e.Code, e.Message) +} + +// newRequest builds a request (id != nil) or notification (id == nil) with the +// given params marshaled to JSON. A nil params value is omitted from the wire +// message. +func newRequest(id *ID, method string, params any) (*Request, error) { + req := &Request{JSONRPC: Version, ID: id, Method: method} + if params != nil { + raw, err := json.Marshal(params) + if err != nil { + return nil, fmt.Errorf("jsonrpc: marshal params for %q: %w", method, err) + } + req.Params = raw + } + return req, nil +} diff --git a/pigo/internal/jsonrpc/transport.go b/pigo/internal/jsonrpc/transport.go new file mode 100644 index 0000000..6baa3e6 --- /dev/null +++ b/pigo/internal/jsonrpc/transport.go @@ -0,0 +1,293 @@ +// This file implements the subprocess JSON-RPC client (US-014/#116). A Client +// launches an external executable, writes requests to its stdin and reads +// newline-delimited JSON-RPC messages from its stdout on a background reader +// goroutine. Requests are correlated to responses by id through a pending-call +// map, so Call is safe for concurrent use: each caller blocks only on its own +// response channel until the reply arrives, the context is cancelled, or the +// child exits. +package jsonrpc + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os/exec" + "sync" + "sync/atomic" + "time" +) + +// closeGrace bounds how long Close waits for a child to exit on its own after +// stdin is closed, before force-killing it. +const closeGrace = 5 * time.Second + +// ErrClosed is returned by Call once the client has been closed or the child +// process has exited. +var ErrClosed = errors.New("jsonrpc: client closed") + +// Client is a JSON-RPC 2.0 client bound to a subprocess over its stdio. +type Client struct { + cmd *exec.Cmd + stdin io.WriteCloser + stdout io.ReadCloser + + writeMu sync.Mutex // serializes writes to the child's stdin + nextID atomic.Int64 + + mu sync.Mutex // guards pending and closed + pending map[string]chan res // id -> waiter + closed bool + closeErr error + + done chan struct{} // closed when the reader goroutine exits +} + +// res carries a decoded response (or a transport-level error) to a waiter. +type res struct { + resp *Response + err error +} + +// Config describes the subprocess to launch. +type Config struct { + // Command is the executable path. + Command string + // Args are the process arguments (excluding the command itself). + Args []string + // Env is the child's environment (os/exec form: "KEY=value"). When nil the + // child inherits the parent environment. + Env []string + // Dir is the child's working directory; empty means the parent's. + Dir string + // Stderr optionally receives the child's stderr (e.g. for logging). When nil + // the child's stderr is discarded. + Stderr io.Writer +} + +// NewClient starts the subprocess and begins reading its stdout. The caller must +// Close the client to terminate the child and release resources. +func NewClient(cfg Config) (*Client, error) { + if cfg.Command == "" { + return nil, errors.New("jsonrpc: empty command") + } + cmd := exec.Command(cfg.Command, cfg.Args...) + cmd.Env = cfg.Env + cmd.Dir = cfg.Dir + cmd.Stderr = cfg.Stderr + + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, fmt.Errorf("jsonrpc: stdin pipe: %w", err) + } + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, fmt.Errorf("jsonrpc: stdout pipe: %w", err) + } + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("jsonrpc: start %q: %w", cfg.Command, err) + } + + c := &Client{ + cmd: cmd, + stdin: stdin, + stdout: stdout, + pending: make(map[string]chan res), + done: make(chan struct{}), + } + go c.readLoop() + return c, nil +} + +// readLoop reads newline-delimited JSON messages from the child's stdout and +// dispatches each response to its waiter. It exits when stdout hits EOF/error, +// failing all outstanding calls. +func (c *Client) readLoop() { + defer close(c.done) + scanner := bufio.NewScanner(c.stdout) + // MCP/plugin payloads (e.g. tool schemas) can be large; raise the line cap. + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + continue + } + var resp Response + if err := json.Unmarshal(line, &resp); err != nil { + // A line we can't parse as a response is skipped (it may be a + // server->client request/notification, which this minimal client + // does not handle). + continue + } + if resp.ID == nil { + continue // notification from server; nothing to correlate + } + c.deliver(resp.ID.String(), res{resp: &resp}) + } + + err := scanner.Err() + if err == nil { + err = io.EOF + } + c.failAll(fmt.Errorf("jsonrpc: reader stopped: %w", err)) +} + +// deliver hands a response to its waiter (if still registered). +func (c *Client) deliver(id string, r res) { + c.mu.Lock() + ch, ok := c.pending[id] + if ok { + delete(c.pending, id) + } + c.mu.Unlock() + if ok { + ch <- r + } +} + +// failAll completes every outstanding call with err and marks the client closed. +func (c *Client) failAll(err error) { + c.mu.Lock() + if c.closeErr == nil { + c.closeErr = err + } + c.closed = true + pending := c.pending + c.pending = make(map[string]chan res) + c.mu.Unlock() + for _, ch := range pending { + ch <- res{err: err} + } +} + +// Call sends a request and waits for the matching response. It returns the +// decoded Result on success, the server Error as error on an error response, or +// a transport/context error. Call is safe for concurrent use. +func (c *Client) Call(ctx context.Context, method string, params any) (json.RawMessage, error) { + id := NumID(c.nextID.Add(1)) + req, err := newRequest(&id, method, params) + if err != nil { + return nil, err + } + + ch := make(chan res, 1) + c.mu.Lock() + if c.closed { + c.mu.Unlock() + return nil, c.closedErr() + } + c.pending[id.String()] = ch + c.mu.Unlock() + + if err := c.write(req); err != nil { + c.mu.Lock() + delete(c.pending, id.String()) + c.mu.Unlock() + return nil, err + } + + select { + case <-ctx.Done(): + c.mu.Lock() + delete(c.pending, id.String()) + c.mu.Unlock() + return nil, ctx.Err() + case r := <-ch: + if r.err != nil { + return nil, r.err + } + if r.resp.Error != nil { + return nil, r.resp.Error + } + return r.resp.Result, nil + } +} + +// Notify sends a notification (no id, no response expected). +func (c *Client) Notify(method string, params any) error { + req, err := newRequest(nil, method, params) + if err != nil { + return err + } + c.mu.Lock() + if c.closed { + c.mu.Unlock() + return c.closedErr() + } + c.mu.Unlock() + return c.write(req) +} + +// write serializes a message and writes it as one newline-terminated line. +// Writes are serialized so concurrent callers don't interleave bytes on stdin. +func (c *Client) write(msg any) error { + data, err := json.Marshal(msg) + if err != nil { + return fmt.Errorf("jsonrpc: marshal request: %w", err) + } + data = append(data, '\n') + c.writeMu.Lock() + defer c.writeMu.Unlock() + if _, err := c.stdin.Write(data); err != nil { + return fmt.Errorf("jsonrpc: write: %w", err) + } + return nil +} + +// closedErr reports why the client is closed, defaulting to ErrClosed. +func (c *Client) closedErr() error { + if c.closeErr != nil { + return c.closeErr + } + return ErrClosed +} + +// Close closes the child's stdin (signalling graceful shutdown), waits briefly +// for the process to exit, and kills it if it does not. It is idempotent. +func (c *Client) Close() error { + c.mu.Lock() + if c.closed { + c.mu.Unlock() + <-c.done + return nil + } + c.closed = true + if c.closeErr == nil { + c.closeErr = ErrClosed + } + c.mu.Unlock() + + // Closing stdin lets a well-behaved child exit on its own. + _ = c.stdin.Close() + + // Wait for the child; kill if it doesn't stop promptly. The grace timer + // bounds Close so a child that ignores stdin EOF and keeps stdout open + // (a hung plugin/server) cannot block us forever. + waitErr := make(chan error, 1) + go func() { waitErr <- c.cmd.Wait() }() + + grace := time.NewTimer(closeGrace) + defer grace.Stop() + + select { + case <-waitErr: + <-c.done + return nil + case <-c.done: + // reader saw EOF; give Wait a brief moment, then force kill. + case <-grace.C: + // child hasn't exited within the grace period; force it. + } + + select { + case <-waitErr: + default: + _ = c.cmd.Process.Kill() + <-waitErr + } + <-c.done + return nil +} diff --git a/pigo/internal/jsonrpc/transport_test.go b/pigo/internal/jsonrpc/transport_test.go new file mode 100644 index 0000000..bd596c8 --- /dev/null +++ b/pigo/internal/jsonrpc/transport_test.go @@ -0,0 +1,223 @@ +package jsonrpc + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "os" + "strings" + "sync" + "testing" + "time" +) + +// TestMain lets this test binary double as a mock JSON-RPC server subprocess. +// When JSONRPC_TEST_SERVER is set the process runs the echo server and exits; +// otherwise it runs the normal test suite. This is the standard Go pattern for +// exercising a subprocess transport without shipping a separate helper binary. +func TestMain(m *testing.M) { + switch os.Getenv("JSONRPC_TEST_SERVER") { + case "echo": + runEchoServer() + return + case "silent": + // Read and discard everything, never reply — used for timeout tests. + sc := bufio.NewScanner(os.Stdin) + sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + for sc.Scan() { + } + return + } + os.Exit(m.Run()) +} + +// runEchoServer replies to each request: method "echo" returns its params, +// method "fail" returns a JSON-RPC error, notifications produce no reply. +func runEchoServer() { + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + w := bufio.NewWriter(os.Stdout) + for scanner.Scan() { + var req Request + if err := json.Unmarshal(scanner.Bytes(), &req); err != nil { + continue + } + if req.ID == nil { + continue // notification: no response + } + var resp Response + resp.JSONRPC = Version + resp.ID = req.ID + switch req.Method { + case "fail": + resp.Error = &Error{Code: -32000, Message: "boom"} + default: + resp.Result = req.Params + if resp.Result == nil { + resp.Result = json.RawMessage(`null`) + } + } + out, _ := json.Marshal(&resp) + out = append(out, '\n') + _, _ = w.Write(out) + _ = w.Flush() + } +} + +// newTestClient starts this test binary as a mock server in the given mode. +func newTestClient(t *testing.T, mode string) *Client { + t.Helper() + exe, err := os.Executable() + if err != nil { + t.Fatalf("os.Executable: %v", err) + } + c, err := NewClient(Config{ + Command: exe, + Env: append(os.Environ(), "JSONRPC_TEST_SERVER="+mode), + }) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + t.Cleanup(func() { _ = c.Close() }) + return c +} + +func TestCallEcho(t *testing.T) { + c := newTestClient(t, "echo") + ctx := context.Background() + + raw, err := c.Call(ctx, "echo", map[string]any{"hello": "world"}) + if err != nil { + t.Fatalf("Call: %v", err) + } + var got map[string]string + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("unmarshal result: %v", err) + } + if got["hello"] != "world" { + t.Fatalf("got %v, want hello=world", got) + } +} + +func TestCallServerError(t *testing.T) { + c := newTestClient(t, "echo") + _, err := c.Call(context.Background(), "fail", nil) + if err == nil { + t.Fatal("expected error, got nil") + } + var rpcErr *Error + if !errors.As(err, &rpcErr) { + t.Fatalf("expected *jsonrpc.Error, got %T: %v", err, err) + } + if rpcErr.Code != -32000 || !strings.Contains(rpcErr.Message, "boom") { + t.Fatalf("unexpected error: %+v", rpcErr) + } +} + +// TestConcurrentCalls verifies responses correlate to the right caller when many +// requests are in flight at once (id-based correlation). +func TestConcurrentCalls(t *testing.T) { + c := newTestClient(t, "echo") + ctx := context.Background() + + const n = 50 + var wg sync.WaitGroup + errs := make([]error, n) + for i := range n { + wg.Add(1) + go func(i int) { + defer wg.Done() + raw, err := c.Call(ctx, "echo", map[string]int{"n": i}) + if err != nil { + errs[i] = err + return + } + var got map[string]int + if err := json.Unmarshal(raw, &got); err != nil { + errs[i] = err + return + } + if got["n"] != i { + errs[i] = fmt.Errorf("call %d got n=%d", i, got["n"]) + } + }(i) + } + wg.Wait() + for i, err := range errs { + if err != nil { + t.Fatalf("call %d: %v", i, err) + } + } +} + +// TestCallContextTimeout verifies Call returns when the context is cancelled and +// the server never replies. +func TestCallContextTimeout(t *testing.T) { + c := newTestClient(t, "silent") + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := c.Call(ctx, "echo", nil) + if err == nil { + t.Fatal("expected timeout error, got nil") + } + if time.Since(start) > 2*time.Second { + t.Fatalf("Call blocked too long: %v", time.Since(start)) + } +} + +func TestNotifyNoResponse(t *testing.T) { + c := newTestClient(t, "echo") + if err := c.Notify("ping", map[string]string{"k": "v"}); err != nil { + t.Fatalf("Notify: %v", err) + } + // A subsequent Call must still work (notification produced no stray reply). + if _, err := c.Call(context.Background(), "echo", nil); err != nil { + t.Fatalf("Call after Notify: %v", err) + } +} + +func TestCallAfterClose(t *testing.T) { + c := newTestClient(t, "echo") + if err := c.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if _, err := c.Call(context.Background(), "echo", nil); err == nil { + t.Fatal("expected error calling closed client") + } +} + +func TestNewClientEmptyCommand(t *testing.T) { + if _, err := NewClient(Config{}); err == nil { + t.Fatal("expected error for empty command") + } +} + +func TestIDRoundTrip(t *testing.T) { + for _, tc := range []struct { + name string + raw string + }{ + {"number", `123`}, + {"string", `"abc"`}, + } { + var id ID + if err := json.Unmarshal([]byte(tc.raw), &id); err != nil { + t.Fatalf("%s: unmarshal: %v", tc.name, err) + } + out, err := json.Marshal(id) + if err != nil { + t.Fatalf("%s: marshal: %v", tc.name, err) + } + if string(out) != tc.raw { + t.Fatalf("%s: round-trip got %s want %s", tc.name, out, tc.raw) + } + } + var bad ID + if err := json.Unmarshal([]byte(`true`), &bad); err == nil { + t.Fatal("expected error for boolean id") + } +} diff --git a/pigo/internal/memory/count_test.go b/pigo/internal/memory/count_test.go new file mode 100644 index 0000000..5f96c70 --- /dev/null +++ b/pigo/internal/memory/count_test.go @@ -0,0 +1,57 @@ +package memory + +import ( + "testing" +) + +// TestCountByScopeEmpty returns an empty map for a fresh store with no entries. +func TestCountByScopeEmpty(t *testing.T) { + st := openTemp(t) + counts, err := st.CountByScope() + if err != nil { + t.Fatalf("CountByScope: %v", err) + } + if len(counts) != 0 { + t.Fatalf("counts = %v, want empty", counts) + } +} + +// TestCountByScopeGroups reconciles files across scopes and asserts the counts +// are grouped per scope. +func TestCountByScopeGroups(t *testing.T) { + st, root, _ := openTempWithRoots(t) + + writeFile(t, root, "g1", "global", "reference", "a.md") + writeFile(t, root, "g2", "global", "notes", "b.md") + writeFile(t, root, "p1", "projects", "proj1", "project", "m.md") + + if _, err := st.Reconcile(); err != nil { + t.Fatalf("Reconcile: %v", err) + } + + counts, err := st.CountByScope() + if err != nil { + t.Fatalf("CountByScope: %v", err) + } + if counts[ScopeGlobal] != 2 { + t.Fatalf("global count = %d, want 2 (counts=%v)", counts[ScopeGlobal], counts) + } + if counts[ScopeProjects] != 1 { + t.Fatalf("projects count = %d, want 1 (counts=%v)", counts[ScopeProjects], counts) + } + if _, ok := counts[ScopeSessions]; ok { + t.Fatalf("sessions should be absent, got %d", counts[ScopeSessions]) + } +} + +// TestCountByScopeNilStore is safe on a nil store and yields an empty map. +func TestCountByScopeNilStore(t *testing.T) { + var st *Store + counts, err := st.CountByScope() + if err != nil { + t.Fatalf("CountByScope on nil: %v", err) + } + if len(counts) != 0 { + t.Fatalf("counts = %v, want empty", counts) + } +} diff --git a/pigo/internal/memory/ftsquery.go b/pigo/internal/memory/ftsquery.go new file mode 100644 index 0000000..b7e5293 --- /dev/null +++ b/pigo/internal/memory/ftsquery.go @@ -0,0 +1,62 @@ +package memory + +import "regexp" + +// ftsTokenRe matches contiguous runs of Unicode letters (incl. CJK), numbers, +// and underscore. Everything else — whitespace, punctuation, FTS5 operator +// characters — is treated as a separator. \p{L} deliberately includes CJK +// letters so queries like "配置文件" tokenize into a single searchable run. +var ftsTokenRe = regexp.MustCompile(`[\p{L}\p{N}_]+`) + +// buildFtsQuery builds an FTS5 MATCH expression from a free-form user query. +// +// FTS5's MATCH grammar has its own operators and special characters +// (`"`, `(`, `)`, `*`, `:`, `^`, `-`, `.`, `{`, `}`). Passing a raw user string +// containing any of these crashes the parser. We tokenize on non-word runs, +// wrap each token in phrase quotes (which turn it into a literal-word search +// that ignores FTS5 special chars), and OR-join. +// +// OR (not AND): AND-join requires EVERY query word to appear in a document, so +// a single descriptive word the user added that is absent from the stored text +// zeroes the whole query even when most tokens match. OR lets BM25 rank by how +// many / how rare the matched tokens are; the caller applies a relative score +// floor to drop common-word-only noise (see Store.Search). +// +// Returns "" when no usable tokens are extracted. Callers treat that as "empty +// query, no results" and send no SQL. +func buildFtsQuery(raw string) string { + matches := ftsTokenRe.FindAllString(raw, -1) + if len(matches) == 0 { + return "" + } + + quoted := make([]string, 0, len(matches)) + for _, tok := range matches { + // Strip any embedded double quotes, then wrap the token as a phrase. + stripped := removeQuotes(tok) + if stripped == "" { + continue + } + quoted = append(quoted, `"`+stripped+`"`) + } + if len(quoted) == 0 { + return "" + } + + out := quoted[0] + for _, q := range quoted[1:] { + out += " OR " + q + } + return out +} + +// removeQuotes strips every double-quote character from s. +func removeQuotes(s string) string { + out := make([]rune, 0, len(s)) + for _, r := range s { + if r != '"' { + out = append(out, r) + } + } + return string(out) +} diff --git a/pigo/internal/memory/ftsquery_test.go b/pigo/internal/memory/ftsquery_test.go new file mode 100644 index 0000000..6279198 --- /dev/null +++ b/pigo/internal/memory/ftsquery_test.go @@ -0,0 +1,55 @@ +package memory + +import "testing" + +func TestBuildFtsQueryMultiWordOrJoin(t *testing.T) { + got := buildFtsQuery("permission deadlock retry") + want := `"permission" OR "deadlock" OR "retry"` + if got != want { + t.Fatalf("buildFtsQuery multi-word: got %q want %q", got, want) + } +} + +func TestBuildFtsQuerySingleToken(t *testing.T) { + if got := buildFtsQuery("checkpoint"); got != `"checkpoint"` { + t.Fatalf("buildFtsQuery single: got %q", got) + } +} + +func TestBuildFtsQueryCJKTokensKept(t *testing.T) { + // CJK letters are \p{L} and must survive tokenization. Whitespace splits + // them into separate tokens; punctuation is a separator. + got := buildFtsQuery("配置文件 端口") + want := `"配置文件" OR "端口"` + if got != want { + t.Fatalf("buildFtsQuery CJK: got %q want %q", got, want) + } +} + +func TestBuildFtsQueryPunctuationStripped(t *testing.T) { + // FTS5 special chars and punctuation become separators; underscores and + // digits are word characters. + got := buildFtsQuery(`port: 5433 (postgres-db) foo_bar`) + want := `"port" OR "5433" OR "postgres" OR "db" OR "foo_bar"` + if got != want { + t.Fatalf("buildFtsQuery punctuation: got %q want %q", got, want) + } +} + +func TestBuildFtsQueryStripsEmbeddedQuotes(t *testing.T) { + // A token can only contain word chars, so a raw double-quote is a + // separator; but guard the strip explicitly. + got := buildFtsQuery(`say "hello"`) + want := `"say" OR "hello"` + if got != want { + t.Fatalf("buildFtsQuery quotes: got %q want %q", got, want) + } +} + +func TestBuildFtsQueryEmptyAndWhitespace(t *testing.T) { + for _, in := range []string{"", " ", "\t\n", "!!! ??? ... ---", "()[]{}"} { + if got := buildFtsQuery(in); got != "" { + t.Fatalf("buildFtsQuery(%q): got %q want empty", in, got) + } + } +} diff --git a/pigo/internal/memory/paths.go b/pigo/internal/memory/paths.go new file mode 100644 index 0000000..bb6e04e --- /dev/null +++ b/pigo/internal/memory/paths.go @@ -0,0 +1,233 @@ +package memory + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "path/filepath" + "strings" + + "gopkg.in/yaml.v3" +) + +// Locator is the parsed identity of a memory file on disk: which scope it +// belongs to, the scope id (project-id / session-id / slug; "" for global), the +// semantic type, and the absolute path of the file itself. +type Locator struct { + Scope Scope + ScopeID string + Type Type + Path string +} + +// knownTypes maps a layout directory segment (or a frontmatter +// metadata.type value) to its Type. Anything not present here resolves to +// TypeFree. +var knownTypes = map[string]Type{ + string(TypeUser): TypeUser, + string(TypeFeedback): TypeFeedback, + string(TypeProject): TypeProject, + string(TypeReference): TypeReference, + string(TypeCheckpoint): TypeCheckpoint, + string(TypeProgress): TypeProgress, + string(TypeNotes): TypeNotes, + string(TypeFree): TypeFree, +} + +// scopeForMarker maps a top-level layout directory name to its Scope. +var scopeForMarker = map[string]Scope{ + "global": ScopeGlobal, + "projects": ScopeProjects, + "sessions": ScopeSessions, +} + +// typeFromDir maps a directory segment to a Type, defaulting to TypeFree +// for unknown segments. +func typeFromDir(seg string) Type { + if t, ok := knownTypes[strings.ToLower(seg)]; ok { + return t + } + return TypeFree +} + +// parsePath parses an absolute path in the mimo memory layout, relative to root, +// into a Locator. Recognized shapes: +// +// /global//*.md -> ScopeGlobal, ScopeID="", Type from +// /projects///*.md -> ScopeProjects, ScopeID=projectId, Type from +// /sessions///*.md -> ScopeSessions, ScopeID=sessionId, Type from +// +// A .md file directly under a scope dir (e.g. /projects//MEMORY.md, or +// /global/MEMORY.md) has no segment and resolves to TypeFree. +// Unknown segments also map to TypeFree. It returns (nil, false) for any +// path outside root or outside the layout. +func parsePath(root, absPath string) (*Locator, bool) { + rel, ok := relComponents(root, absPath) + if !ok || len(rel) == 0 { + return nil, false + } + + scope, ok := scopeForMarker[rel[0]] + if !ok { + return nil, false + } + + // rest is everything below the top-level scope marker. + rest := rel[1:] + + var scopeID string + switch scope { + case ScopeGlobal: + // rest = [/].md + case ScopeProjects, ScopeSessions: + // rest = /[/].md + if len(rest) == 0 { + return nil, false + } + scopeID = rest[0] + rest = rest[1:] + default: + return nil, false + } + + // A file component is mandatory and must be a Markdown file. + if len(rest) == 0 || !isMarkdown(rest[len(rest)-1]) { + return nil, false + } + + typ := TypeFree + if len(rest) >= 2 { + // rest = /.../.md — the first segment names the type. + typ = typeFromDir(rest[0]) + } + + return &Locator{ + Scope: scope, + ScopeID: scopeID, + Type: typ, + Path: filepath.Clean(absPath), + }, true +} + +// parseCcPath parses a Claude Code layout path relative to ccBase: +// +// //memory/**/*.md -> ScopeCC, ScopeID=, Type=TypeFree +// +// The real type is derived later from the file's YAML frontmatter via +// parseCcFrontmatterType. It returns (nil, false) for paths outside the layout. +func parseCcPath(ccBase, absPath string) (*Locator, bool) { + rel, ok := relComponents(ccBase, absPath) + if !ok { + return nil, false + } + // Need at least /memory/.md. + if len(rel) < 3 || rel[1] != "memory" { + return nil, false + } + if !isMarkdown(rel[len(rel)-1]) { + return nil, false + } + return &Locator{ + Scope: ScopeCC, + ScopeID: rel[0], + Type: TypeFree, + Path: filepath.Clean(absPath), + }, true +} + +// ccFrontmatter is the minimal shape read from a Claude Code memory file's +// leading YAML frontmatter: either a nested metadata.type or a top-level type. +type ccFrontmatter struct { + Type string `yaml:"type"` + Metadata struct { + Type string `yaml:"type"` + } `yaml:"metadata"` +} + +// parseCcFrontmatterType reads the semantic type from a leading YAML frontmatter +// block (a "---" fence at the very start of body). It prefers metadata.type, +// falling back to a top-level type. Absent, malformed, or unrecognized values +// yield TypeFree. +func parseCcFrontmatterType(body string) Type { + block, ok := frontmatterBlock(body) + if !ok { + return TypeFree + } + var fm ccFrontmatter + if err := yaml.Unmarshal([]byte(block), &fm); err != nil { + return TypeFree + } + candidate := fm.Metadata.Type + if candidate == "" { + candidate = fm.Type + } + if t, ok := knownTypes[strings.ToLower(strings.TrimSpace(candidate))]; ok { + return t + } + return TypeFree +} + +// resolveProjectId derives a stable project id from an absolute repository path: +// the first 12 hex characters of sha256(absRepoPath). It is deterministic and is +// used as the scope_id for the projects scope. +func resolveProjectId(absRepoPath string) string { + sum := sha256.Sum256([]byte(absRepoPath)) + return hex.EncodeToString(sum[:])[:12] +} + +// assertSafeComponent rejects a caller-supplied path component that would escape +// the memory root: any ".." segment or a leading "/" (absolute path). Empty +// components are also rejected. The write path uses this to sanitize +// scope_id/type/filename before joining them under root. +func assertSafeComponent(name string) error { + if name == "" { + return fmt.Errorf("memory: empty path component") + } + if strings.HasPrefix(name, "/") { + return fmt.Errorf("memory: unsafe path component %q: leading %q", name, "/") + } + for _, seg := range strings.Split(filepath.ToSlash(name), "/") { + if seg == ".." { + return fmt.Errorf("memory: unsafe path component %q: %q segment", name, "..") + } + } + return nil +} + +// relComponents returns absPath expressed relative to base, split into +// non-empty components. ok is false when absPath is not located under base. +func relComponents(base, absPath string) (parts []string, ok bool) { + rel, err := filepath.Rel(filepath.Clean(base), filepath.Clean(absPath)) + if err != nil { + return nil, false + } + rel = filepath.ToSlash(rel) + if rel == "." || rel == "" { + return nil, false + } + // Escaping base (e.g. "../foo") means the path is outside the layout. + if rel == ".." || strings.HasPrefix(rel, "../") { + return nil, false + } + return strings.Split(rel, "/"), true +} + +// isMarkdown reports whether name has a .md extension (case-insensitive). +func isMarkdown(name string) bool { + return strings.EqualFold(filepath.Ext(name), ".md") +} + +// frontmatterBlock returns the raw YAML between a leading "---" fence and the +// next "---" line. ok is false when body has no opening fence at its very start. +func frontmatterBlock(body string) (string, bool) { + rest := strings.ReplaceAll(body, "\r\n", "\n") + if !strings.HasPrefix(rest, "---\n") { + return "", false + } + rest = strings.TrimPrefix(rest, "---\n") + end := strings.Index(rest, "\n---") + if end < 0 { + return "", false + } + return rest[:end], true +} diff --git a/pigo/internal/memory/paths_test.go b/pigo/internal/memory/paths_test.go new file mode 100644 index 0000000..2918955 --- /dev/null +++ b/pigo/internal/memory/paths_test.go @@ -0,0 +1,203 @@ +package memory + +import ( + "path/filepath" + "testing" +) + +func TestParsePathScopesAndTypes(t *testing.T) { + root := "/mem/root" + + cases := []struct { + name string + path string + scope Scope + scopeID string + typ Type + }{ + { + name: "global with type dir", + path: filepath.Join(root, "global", "user", "profile.md"), + scope: ScopeGlobal, scopeID: "", typ: TypeUser, + }, + { + name: "global unknown type dir -> free", + path: filepath.Join(root, "global", "whatever", "x.md"), + scope: ScopeGlobal, scopeID: "", typ: TypeFree, + }, + { + name: "global file directly under scope -> free", + path: filepath.Join(root, "global", "MEMORY.md"), + scope: ScopeGlobal, scopeID: "", typ: TypeFree, + }, + { + name: "projects with type dir", + path: filepath.Join(root, "projects", "abc123", "checkpoint", "c1.md"), + scope: ScopeProjects, scopeID: "abc123", typ: TypeCheckpoint, + }, + { + name: "projects file directly under id -> free", + path: filepath.Join(root, "projects", "abc123", "MEMORY.md"), + scope: ScopeProjects, scopeID: "abc123", typ: TypeFree, + }, + { + name: "sessions with type dir", + path: filepath.Join(root, "sessions", "sess-1", "notes", "n.md"), + scope: ScopeSessions, scopeID: "sess-1", typ: TypeNotes, + }, + { + name: "sessions progress type", + path: filepath.Join(root, "sessions", "sess-1", "progress", "p.md"), + scope: ScopeSessions, scopeID: "sess-1", typ: TypeProgress, + }, + { + name: "nested file under type dir keeps type", + path: filepath.Join(root, "projects", "abc123", "reference", "sub", "r.md"), + scope: ScopeProjects, scopeID: "abc123", typ: TypeReference, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + loc, ok := parsePath(root, tc.path) + if !ok { + t.Fatalf("parsePath(%q) returned ok=false", tc.path) + } + if loc.Scope != tc.scope { + t.Errorf("scope = %q, want %q", loc.Scope, tc.scope) + } + if loc.ScopeID != tc.scopeID { + t.Errorf("scopeID = %q, want %q", loc.ScopeID, tc.scopeID) + } + if loc.Type != tc.typ { + t.Errorf("type = %q, want %q", loc.Type, tc.typ) + } + if loc.Path != filepath.Clean(tc.path) { + t.Errorf("path = %q, want %q", loc.Path, filepath.Clean(tc.path)) + } + }) + } +} + +func TestParsePathOutsideLayout(t *testing.T) { + root := "/mem/root" + bad := []string{ + "/other/place/x.md", // outside root + filepath.Join(root, "unknownscope", "x.md"), // not a layout scope + filepath.Join(root, "global"), // no file component + filepath.Join(root, "projects", "abc123"), // scope id dir, no file + filepath.Join(root, "global", "user", "x.txt"), // not markdown + } + for _, p := range bad { + if loc, ok := parsePath(root, p); ok { + t.Errorf("parsePath(%q) = %+v, ok=true; want ok=false", p, loc) + } + } +} + +func TestParseCcPath(t *testing.T) { + base := "/home/u/.claude/projects" + + loc, ok := parseCcPath(base, filepath.Join(base, "my-slug", "memory", "some", "note.md")) + if !ok { + t.Fatalf("parseCcPath returned ok=false") + } + if loc.Scope != ScopeCC { + t.Errorf("scope = %q, want %q", loc.Scope, ScopeCC) + } + if loc.ScopeID != "my-slug" { + t.Errorf("scopeID = %q, want %q", loc.ScopeID, "my-slug") + } + if loc.Type != TypeFree { + t.Errorf("type = %q, want %q", loc.Type, TypeFree) + } + + bad := []string{ + filepath.Join(base, "my-slug", "note.md"), // no memory segment + filepath.Join(base, "my-slug", "notmemory", "n.md"), // wrong segment + filepath.Join(base, "my-slug", "memory", "note.txt"), // not markdown + "/elsewhere/x/memory/n.md", // outside base + } + for _, p := range bad { + if loc, ok := parseCcPath(base, p); ok { + t.Errorf("parseCcPath(%q) = %+v, ok=true; want ok=false", p, loc) + } + } +} + +func TestParseCcFrontmatterType(t *testing.T) { + cases := []struct { + name string + body string + want Type + }{ + { + name: "nested metadata.type present", + body: "---\nmetadata:\n type: reference\n other: x\n---\nbody here\n", + want: TypeReference, + }, + { + name: "top-level type present", + body: "---\ntype: feedback\n---\nbody\n", + want: TypeFeedback, + }, + { + name: "metadata.type wins over top-level", + body: "---\ntype: free\nmetadata:\n type: project\n---\n", + want: TypeProject, + }, + { + name: "absent frontmatter", + body: "no frontmatter here\n", + want: TypeFree, + }, + { + name: "empty frontmatter", + body: "---\n---\nbody\n", + want: TypeFree, + }, + { + name: "unknown type value", + body: "---\nmetadata:\n type: bogus\n---\n", + want: TypeFree, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := parseCcFrontmatterType(tc.body); got != tc.want { + t.Errorf("parseCcFrontmatterType = %q, want %q", got, tc.want) + } + }) + } +} + +func TestResolveProjectId(t *testing.T) { + const p = "/Users/dev/repo" + a := resolveProjectId(p) + b := resolveProjectId(p) + if a != b { + t.Errorf("resolveProjectId not stable: %q != %q", a, b) + } + if len(a) != 12 { + t.Errorf("resolveProjectId length = %d, want 12", len(a)) + } + if resolveProjectId("/Users/dev/other") == a { + t.Errorf("resolveProjectId collided for distinct paths") + } +} + +func TestAssertSafeComponent(t *testing.T) { + ok := []string{"user", "abc123", "sub/dir/file.md", "checkpoint"} + for _, s := range ok { + if err := assertSafeComponent(s); err != nil { + t.Errorf("assertSafeComponent(%q) = %v, want nil", s, err) + } + } + + bad := []string{"", "..", "../etc", "a/../b", "/etc/passwd", "sub/../../x"} + for _, s := range bad { + if err := assertSafeComponent(s); err == nil { + t.Errorf("assertSafeComponent(%q) = nil, want error", s) + } + } +} diff --git a/pigo/internal/memory/reconcile.go b/pigo/internal/memory/reconcile.go new file mode 100644 index 0000000..e291a1b --- /dev/null +++ b/pigo/internal/memory/reconcile.go @@ -0,0 +1,235 @@ +package memory + +import ( + "fmt" + "os" + "path/filepath" + "time" +) + +// Result reports what a Reconcile pass changed: Indexed counts rows inserted or +// updated (new or changed files), Pruned counts rows deleted because their file +// no longer exists on disk. +type Result struct { + Indexed int + Pruned int +} + +// walkMemoryDir recursively collects every *.md file under root. A missing root +// (ENOENT) yields an empty slice and no error, so reconcile is safe to run +// before the memory directory has been created. +func walkMemoryDir(root string) ([]string, error) { + var out []string + var recurse func(dir string) error + recurse = func(dir string) error { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + for _, entry := range entries { + full := filepath.Join(dir, entry.Name()) + if entry.IsDir() { + if err := recurse(full); err != nil { + return err + } + } else if entry.Type().IsRegular() && isMarkdown(entry.Name()) { + out = append(out, full) + } + } + return nil + } + if err := recurse(root); err != nil { + return nil, err + } + return out, nil +} + +// walkCcRoot collects every //memory/**/*.md file. A missing base +// (ENOENT) yields an empty slice; slugs without a memory subdirectory are +// silently skipped. +func walkCcRoot(base string) ([]string, error) { + slugs, err := os.ReadDir(base) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, err + } + var out []string + for _, entry := range slugs { + if !entry.IsDir() { + continue + } + memoryDir := filepath.Join(base, entry.Name(), "memory") + info, err := os.Stat(memoryDir) + if err != nil || !info.IsDir() { + continue + } + files, err := walkMemoryDir(memoryDir) + if err != nil { + return nil, err + } + out = append(out, files...) + } + return out, nil +} + +// Reconcile performs a lazy sync between the memory files on disk and the +// memory_index table. It walks both the mimo root and (when configured) the cc +// base, prunes rows whose file no longer exists, and indexes new or changed +// files. Unchanged files are skipped via a size-mtime fingerprint. The FTS +// index is kept consistent by the memory_ai/ad/au triggers, so only +// memory_index is touched here. +func (s *Store) Reconcile() (Result, error) { + var res Result + + // Collect disk paths from BOTH roots BEFORE pruning. Pruning per-root would + // wrongly wipe the other root's rows, because each walk's set is missing the + // other root's paths. + mimoFiles, err := walkMemoryDir(s.root) + if err != nil { + return res, fmt.Errorf("memory: walk mimo root %q: %w", s.root, err) + } + var ccFiles []string + if s.ccBase != "" { + ccFiles, err = walkCcRoot(s.ccBase) + if err != nil { + return res, fmt.Errorf("memory: walk cc base %q: %w", s.ccBase, err) + } + } + + diskPaths := make(map[string]struct{}, len(mimoFiles)+len(ccFiles)) + for _, p := range mimoFiles { + diskPaths[filepath.Clean(p)] = struct{}{} + } + for _, p := range ccFiles { + diskPaths[filepath.Clean(p)] = struct{}{} + } + + // Load existing {path -> fingerprint} from memory_index. + existing, err := s.loadFingerprints() + if err != nil { + return res, err + } + + // PRUNE: delete rows whose path is no longer on disk. + for p := range existing { + if _, ok := diskPaths[p]; ok { + continue + } + if _, err := s.db.Exec(`DELETE FROM memory_index WHERE path = ?`, p); err != nil { + return res, fmt.Errorf("memory: prune %q: %w", p, err) + } + res.Pruned++ + } + + // INDEX: mimo files use parsePath and keep loc.Type. + for _, p := range mimoFiles { + loc, ok := parsePath(s.root, p) + if !ok { + continue + } + updated, err := s.indexFile(loc, loc.Type, false, existing[filepath.Clean(p)]) + if err != nil { + return res, err + } + if updated { + res.Indexed++ + } + } + + // INDEX: cc files use parseCcPath; final type is derived from frontmatter. + for _, p := range ccFiles { + loc, ok := parseCcPath(s.ccBase, p) + if !ok { + continue + } + updated, err := s.indexFile(loc, loc.Type, true, existing[filepath.Clean(p)]) + if err != nil { + return res, err + } + if updated { + res.Indexed++ + } + } + + return res, nil +} + +// loadFingerprints returns the current {path -> fingerprint} map from +// memory_index. +func (s *Store) loadFingerprints() (map[string]string, error) { + rows, err := s.db.Query(`SELECT path, fingerprint FROM memory_index`) + if err != nil { + return nil, fmt.Errorf("memory: load fingerprints: %w", err) + } + defer rows.Close() + + out := make(map[string]string) + for rows.Next() { + var path, fp string + if err := rows.Scan(&path, &fp); err != nil { + return nil, fmt.Errorf("memory: scan fingerprint: %w", err) + } + out[path] = fp + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("memory: iterate fingerprints: %w", err) + } + return out, nil +} + +// indexFile stats loc.Path, computes its size-mtime fingerprint, and upserts the +// row when the fingerprint differs from oldFingerprint. It returns true when a +// row was inserted or updated. A file that vanished between the walk and the +// stat (ENOENT) is silently skipped. For cc files (isCc) the semantic type is +// derived from the file's YAML frontmatter, falling back to defaultType. +func (s *Store) indexFile(loc *Locator, defaultType Type, isCc bool, oldFingerprint string) (bool, error) { + info, err := os.Stat(loc.Path) + if err != nil { + if os.IsNotExist(err) { + return false, nil + } + return false, fmt.Errorf("memory: stat %q: %w", loc.Path, err) + } + + fingerprint := fmt.Sprintf("%d-%d", info.Size(), info.ModTime().UnixNano()) + if oldFingerprint == fingerprint { + return false, nil // hit: unchanged file + } + + raw, err := os.ReadFile(loc.Path) + if err != nil { + if os.IsNotExist(err) { + return false, nil + } + return false, fmt.Errorf("memory: read %q: %w", loc.Path, err) + } + body := string(raw) + + finalType := defaultType + if isCc { + finalType = parseCcFrontmatterType(body) + } + + now := time.Now().UnixNano() + const upsert = ` +INSERT INTO memory_index (path, scope, scope_id, type, body, fingerprint, last_indexed_at) +VALUES (?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(path) DO UPDATE SET + scope = excluded.scope, + scope_id = excluded.scope_id, + type = excluded.type, + body = excluded.body, + fingerprint = excluded.fingerprint, + last_indexed_at = excluded.last_indexed_at` + if _, err := s.db.Exec(upsert, + loc.Path, string(loc.Scope), loc.ScopeID, string(finalType), body, fingerprint, now, + ); err != nil { + return false, fmt.Errorf("memory: upsert %q: %w", loc.Path, err) + } + return true, nil +} diff --git a/pigo/internal/memory/reconcile_test.go b/pigo/internal/memory/reconcile_test.go new file mode 100644 index 0000000..430d11a --- /dev/null +++ b/pigo/internal/memory/reconcile_test.go @@ -0,0 +1,256 @@ +package memory + +import ( + "os" + "path/filepath" + "testing" + "time" +) + +// openTempWithRoots opens a Store backed by a temp-file DB with explicit mimo +// root and cc base directories (both created on disk). +func openTempWithRoots(t *testing.T) (st *Store, root, ccBase string) { + t.Helper() + base := t.TempDir() + root = filepath.Join(base, "mimo") + ccBase = filepath.Join(base, "cc") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatalf("mkdir root: %v", err) + } + if err := os.MkdirAll(ccBase, 0o755); err != nil { + t.Fatalf("mkdir ccBase: %v", err) + } + dbPath := filepath.Join(base, "sub", "memory.db") + st, err := Open(dbPath, root, ccBase) + if err != nil { + t.Fatalf("Open: %v", err) + } + t.Cleanup(func() { st.Close() }) + return st, root, ccBase +} + +// writeFile writes body to a mimo path built from segments under root, creating +// parent directories. It returns the absolute path (cleaned). +func writeFile(t *testing.T, base string, body string, segs ...string) string { + t.Helper() + full := filepath.Join(append([]string{base}, segs...)...) + if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil { + t.Fatalf("mkdir for %q: %v", full, err) + } + if err := os.WriteFile(full, []byte(body), 0o644); err != nil { + t.Fatalf("write %q: %v", full, err) + } + return filepath.Clean(full) +} + +// rowFor returns (scope, scopeId, type, fingerprint, body, found) for a path. +func rowFor(t *testing.T, st *Store, path string) (scope, scopeID, typ, fp, body string, found bool) { + t.Helper() + err := st.DB().QueryRow( + `SELECT scope, scope_id, type, fingerprint, body FROM memory_index WHERE path = ?`, path, + ).Scan(&scope, &scopeID, &typ, &fp, &body) + if err != nil { + return "", "", "", "", "", false + } + return scope, scopeID, typ, fp, body, true +} + +func countRows(t *testing.T, st *Store) int { + t.Helper() + var n int + if err := st.DB().QueryRow(`SELECT count(*) FROM memory_index`).Scan(&n); err != nil { + t.Fatalf("count rows: %v", err) + } + return n +} + +// TestReconcileNewFileIndexed indexes a fresh mimo file and records its locator. +func TestReconcileNewFileIndexed(t *testing.T) { + st, root, _ := openTempWithRoots(t) + + p := writeFile(t, root, "hello world", "global", "reference", "note.md") + + res, err := st.Reconcile() + if err != nil { + t.Fatalf("Reconcile: %v", err) + } + if res.Indexed != 1 || res.Pruned != 0 { + t.Fatalf("Reconcile result = %+v, want {Indexed:1 Pruned:0}", res) + } + + scope, scopeID, typ, fp, body, found := rowFor(t, st, p) + if !found { + t.Fatalf("row for %q not found", p) + } + if scope != string(ScopeGlobal) || scopeID != "" || typ != string(TypeReference) { + t.Fatalf("row = scope=%q scope_id=%q type=%q, want global/''/reference", scope, scopeID, typ) + } + if body != "hello world" { + t.Fatalf("body = %q, want %q", body, "hello world") + } + if fp == "" { + t.Fatalf("fingerprint empty") + } +} + +// TestReconcileUnchangedFileHit re-runs reconcile with no changes and expects a +// fingerprint hit (no re-index). +func TestReconcileUnchangedFileHit(t *testing.T) { + st, root, _ := openTempWithRoots(t) + writeFile(t, root, "stable content", "global", "notes", "a.md") + + if res, err := st.Reconcile(); err != nil || res.Indexed != 1 { + t.Fatalf("first Reconcile = %+v, err=%v, want Indexed:1", res, err) + } + res, err := st.Reconcile() + if err != nil { + t.Fatalf("second Reconcile: %v", err) + } + if res.Indexed != 0 || res.Pruned != 0 { + t.Fatalf("second Reconcile = %+v, want {Indexed:0 Pruned:0} (fingerprint hit)", res) + } +} + +// TestReconcileChangedFileReindexed bumps a file's size and mtime and expects a +// re-index. +func TestReconcileChangedFileReindexed(t *testing.T) { + st, root, _ := openTempWithRoots(t) + p := writeFile(t, root, "v1", "global", "notes", "a.md") + + if _, err := st.Reconcile(); err != nil { + t.Fatalf("first Reconcile: %v", err) + } + _, _, _, fp1, _, _ := rowFor(t, st, p) + + // Rewrite with different size and force a later mtime to guarantee the + // fingerprint changes regardless of filesystem timestamp resolution. + if err := os.WriteFile(p, []byte("v2 longer body"), 0o644); err != nil { + t.Fatalf("rewrite: %v", err) + } + future := time.Now().Add(2 * time.Second) + if err := os.Chtimes(p, future, future); err != nil { + t.Fatalf("chtimes: %v", err) + } + + res, err := st.Reconcile() + if err != nil { + t.Fatalf("second Reconcile: %v", err) + } + if res.Indexed != 1 || res.Pruned != 0 { + t.Fatalf("second Reconcile = %+v, want {Indexed:1 Pruned:0}", res) + } + _, _, _, fp2, body, _ := rowFor(t, st, p) + if fp1 == fp2 { + t.Fatalf("fingerprint unchanged after edit: %q", fp2) + } + if body != "v2 longer body" { + t.Fatalf("body = %q, want re-indexed content", body) + } +} + +// TestReconcileDeletedFilePruned removes a file and expects its row pruned. +func TestReconcileDeletedFilePruned(t *testing.T) { + st, root, _ := openTempWithRoots(t) + p := writeFile(t, root, "temp", "global", "notes", "gone.md") + + if _, err := st.Reconcile(); err != nil { + t.Fatalf("first Reconcile: %v", err) + } + if countRows(t, st) != 1 { + t.Fatalf("want 1 row after index") + } + + if err := os.Remove(p); err != nil { + t.Fatalf("remove: %v", err) + } + res, err := st.Reconcile() + if err != nil { + t.Fatalf("second Reconcile: %v", err) + } + if res.Indexed != 0 || res.Pruned != 1 { + t.Fatalf("second Reconcile = %+v, want {Indexed:0 Pruned:1}", res) + } + if _, _, _, _, _, found := rowFor(t, st, p); found { + t.Fatalf("row for deleted file still present") + } +} + +// TestReconcileCcFrontmatterType indexes a cc-root file and derives its type +// from YAML frontmatter. +func TestReconcileCcFrontmatterType(t *testing.T) { + st, _, ccBase := openTempWithRoots(t) + + const body = "---\nmetadata:\n type: checkpoint\n---\ncc body text" + // //memory/**/*.md + p := writeFile(t, ccBase, body, "my-project", "memory", "sub", "cp.md") + + res, err := st.Reconcile() + if err != nil { + t.Fatalf("Reconcile: %v", err) + } + if res.Indexed != 1 || res.Pruned != 0 { + t.Fatalf("Reconcile = %+v, want {Indexed:1 Pruned:0}", res) + } + scope, scopeID, typ, _, _, found := rowFor(t, st, p) + if !found { + t.Fatalf("cc row not found") + } + if scope != string(ScopeCC) || scopeID != "my-project" || typ != string(TypeCheckpoint) { + t.Fatalf("cc row = scope=%q scope_id=%q type=%q, want cc/my-project/checkpoint", scope, scopeID, typ) + } +} + +// TestReconcileBothRootsNoCrossPrune verifies that indexing both roots does not +// prune the other root's rows (the reconcile.ts correctness note). +func TestReconcileBothRootsNoCrossPrune(t *testing.T) { + st, root, ccBase := openTempWithRoots(t) + + mimoP := writeFile(t, root, "mimo body", "projects", "proj1", "project", "m.md") + ccP := writeFile(t, ccBase, "cc body", "slug1", "memory", "c.md") + + res, err := st.Reconcile() + if err != nil { + t.Fatalf("Reconcile: %v", err) + } + if res.Indexed != 2 || res.Pruned != 0 { + t.Fatalf("Reconcile = %+v, want {Indexed:2 Pruned:0}", res) + } + if _, _, _, _, _, ok := rowFor(t, st, mimoP); !ok { + t.Fatalf("mimo row missing") + } + if _, _, _, _, _, ok := rowFor(t, st, ccP); !ok { + t.Fatalf("cc row missing") + } + + // A second no-op reconcile must not prune either row. + res, err = st.Reconcile() + if err != nil { + t.Fatalf("second Reconcile: %v", err) + } + if res.Pruned != 0 { + t.Fatalf("second Reconcile pruned %d, want 0", res.Pruned) + } + if countRows(t, st) != 2 { + t.Fatalf("row count = %d, want 2", countRows(t, st)) + } +} + +// TestReconcileMissingRoots verifies reconcile is a no-op when roots do not +// exist yet (ENOENT tolerated). +func TestReconcileMissingRoots(t *testing.T) { + base := t.TempDir() + dbPath := filepath.Join(base, "memory.db") + st, err := Open(dbPath, filepath.Join(base, "nope"), filepath.Join(base, "nocc")) + if err != nil { + t.Fatalf("Open: %v", err) + } + t.Cleanup(func() { st.Close() }) + + res, err := st.Reconcile() + if err != nil { + t.Fatalf("Reconcile on missing roots: %v", err) + } + if res.Indexed != 0 || res.Pruned != 0 { + t.Fatalf("Reconcile = %+v, want zero", res) + } +} diff --git a/pigo/internal/memory/schema.go b/pigo/internal/memory/schema.go new file mode 100644 index 0000000..5cef9c4 --- /dev/null +++ b/pigo/internal/memory/schema.go @@ -0,0 +1,66 @@ +package memory + +// Scope identifies the memory dimension a file belongs to. +type Scope string + +// Recognized scopes. +const ( + ScopeGlobal Scope = "global" + ScopeProjects Scope = "projects" + ScopeSessions Scope = "sessions" + ScopeCC Scope = "cc" +) + +// Type identifies the semantic kind of a memory file. +type Type string + +// Recognized types. +const ( + TypeUser Type = "user" + TypeFeedback Type = "feedback" + TypeProject Type = "project" + TypeReference Type = "reference" + TypeCheckpoint Type = "checkpoint" + TypeProgress Type = "progress" + TypeNotes Type = "notes" + TypeFree Type = "free" +) + +// schemaDDL is the idempotent set of DDL statements that create the memory +// storage schema: the content table, its secondary indexes, the FTS5 virtual +// table (external-content mode), and the three sync triggers that keep the FTS +// index consistent with the content table. All statements use IF NOT EXISTS so +// running the migration repeatedly is safe. +const schemaDDL = ` +CREATE TABLE IF NOT EXISTS memory_index ( + id INTEGER PRIMARY KEY, + path TEXT NOT NULL UNIQUE, + scope TEXT NOT NULL, + scope_id TEXT NOT NULL DEFAULT '', + type TEXT NOT NULL, + body TEXT NOT NULL, + fingerprint TEXT NOT NULL, + last_indexed_at INTEGER NOT NULL +); + +CREATE INDEX IF NOT EXISTS memory_index_scope_idx ON memory_index (scope, scope_id); +CREATE INDEX IF NOT EXISTS memory_index_type_idx ON memory_index (type); + +CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5( + body, content='memory_index', content_rowid='id', + tokenize='unicode61 remove_diacritics 1' +); + +CREATE TRIGGER IF NOT EXISTS memory_ai AFTER INSERT ON memory_index BEGIN + INSERT INTO memory_fts(rowid, body) VALUES (new.id, new.body); +END; + +CREATE TRIGGER IF NOT EXISTS memory_ad AFTER DELETE ON memory_index BEGIN + INSERT INTO memory_fts(memory_fts, rowid, body) VALUES('delete', old.id, old.body); +END; + +CREATE TRIGGER IF NOT EXISTS memory_au AFTER UPDATE ON memory_index BEGIN + INSERT INTO memory_fts(memory_fts, rowid, body) VALUES('delete', old.id, old.body); + INSERT INTO memory_fts(rowid, body) VALUES (new.id, new.body); +END; +` diff --git a/pigo/internal/memory/search.go b/pigo/internal/memory/search.go new file mode 100644 index 0000000..ba54361 --- /dev/null +++ b/pigo/internal/memory/search.go @@ -0,0 +1,173 @@ +package memory + +import ( + "database/sql" + "fmt" + "strings" +) + +// SearchResult is a single ranked hit from Store.Search. Score is normalized so +// that higher = better (the raw FTS5 bm25 value, where lower is better, is +// negated). +type SearchResult struct { + Path string + Snippet string + Score float64 + Scope Scope + ScopeID string + Type Type +} + +// SearchOptions tunes a Store.Search call. All fields are optional; the zero +// value performs an unfiltered search with default limit and score floor. +type SearchOptions struct { + // Scope, ScopeID and Type, when non-empty, restrict results to rows whose + // corresponding memory_index column matches exactly. + Scope string + ScopeID string + Type string + // Limit caps the number of returned results; <=0 means the default of 10. + Limit int + // ReconcileFirst runs a lazy Store.Reconcile before searching so that + // off-tool writes are picked up. Its counts are ignored; hard errors + // propagate. + ReconcileFirst bool + // ScoreFloor is the relative floor ratio: trailing rows scoring below + // topScore*ScoreFloor are dropped. <=0 keeps all matches when 0, but the + // default of 0.15 is used when the field is left at its zero value; pass a + // negative value to explicitly disable the floor. + ScoreFloor float64 +} + +// defaultSearchLimit is the result count used when SearchOptions.Limit <= 0. +const defaultSearchLimit = 10 + +// maxFetchLimit caps the over-fetch used to feed the relative score floor. +const maxFetchLimit = 50 + +// defaultScoreFloor is the relative floor applied when SearchOptions.ScoreFloor +// is left at its zero value. +const defaultScoreFloor = 0.15 + +// Search runs a BM25 full-text query over the indexed memory bodies. +// +// The free-form query is tokenized and OR-joined by buildFtsQuery; an empty +// token set returns (nil, nil) without touching SQL. Results are ranked by +// BM25 (converted to higher = better), over-fetched 3x (capped at 50) so a +// relative score floor can trim common-word-only noise, then sliced to the +// requested limit. Optional scope/scope_id/type filters restrict the corpus. +func (s *Store) Search(query string, opts SearchOptions) ([]SearchResult, error) { + if opts.ReconcileFirst { + if _, err := s.Reconcile(); err != nil { + return nil, fmt.Errorf("memory: reconcile before search: %w", err) + } + } + + match := buildFtsQuery(query) + if match == "" { + // No usable tokens: treat as empty query with no results, send no SQL. + return nil, nil + } + + limit := opts.Limit + if limit <= 0 { + limit = defaultSearchLimit + } + fetchLimit := limit * 3 + if fetchLimit > maxFetchLimit { + fetchLimit = maxFetchLimit + } + + // The FTS5 table `memory_fts` is external-content over `memory_index`, so + // filter columns (scope/scope_id/type) and the display path live on + // memory_index and the join is memory_index.id = memory_fts.rowid. snippet() + // and bm25() operate on the FTS table; body is FTS column 0. + var sb strings.Builder + sb.WriteString(` +SELECT mi.path, mi.scope, mi.scope_id, mi.type, + snippet(memory_fts, 0, '<<', '>>', '...', 32) AS snippet, + bm25(memory_fts) AS score +FROM memory_fts +JOIN memory_index mi ON mi.id = memory_fts.rowid +WHERE memory_fts MATCH ?`) + + // MATCH parameter is always first; filter params follow in order. + args := []any{match} + if opts.Scope != "" { + sb.WriteString(" AND mi.scope = ?") + args = append(args, opts.Scope) + } + if opts.ScopeID != "" { + sb.WriteString(" AND mi.scope_id = ?") + args = append(args, opts.ScopeID) + } + if opts.Type != "" { + sb.WriteString(" AND mi.type = ?") + args = append(args, opts.Type) + } + // bm25(): lower = better, so ascending order puts the best hit first. + sb.WriteString(" ORDER BY score ASC LIMIT ?") + args = append(args, fetchLimit) + + rows, err := s.db.Query(sb.String(), args...) + if err != nil { + return nil, fmt.Errorf("memory: search query: %w", err) + } + defer rows.Close() + + var results []SearchResult + for rows.Next() { + var ( + path, scope, scopeID, typ string + snippet sql.NullString + bm25 float64 + ) + if err := rows.Scan(&path, &scope, &scopeID, &typ, &snippet, &bm25); err != nil { + return nil, fmt.Errorf("memory: scan search row: %w", err) + } + results = append(results, SearchResult{ + Path: path, + Snippet: snippet.String, + Score: -bm25, // convert to higher = better + Scope: Scope(scope), + ScopeID: scopeID, + Type: Type(typ), + }) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("memory: iterate search rows: %w", err) + } + if len(results) == 0 { + return nil, nil + } + + // Relative score floor. BM25 magnitudes are corpus-size dependent (in a + // tiny corpus every score collapses toward 0 due to low IDF), so an + // absolute floor would wrongly wipe real hits. We keep results scoring at + // least topScore*floor. The #1 result is ALWAYS kept — a match is a match + // even when BM25 can't discriminate. Default 0.15; a negative floor + // disables the trimming entirely. + floor := opts.ScoreFloor + if floor == 0 { + floor = defaultScoreFloor + } + + // Rows come back ORDER BY score ASC (best first after negation), so + // results[0] is the top hit. + if floor > 0 { + topScore := results[0].Score + cutoff := topScore * floor + kept := results[:1] + for _, r := range results[1:] { + if r.Score >= cutoff { + kept = append(kept, r) + } + } + results = kept + } + + if len(results) > limit { + results = results[:limit] + } + return results, nil +} diff --git a/pigo/internal/memory/search_test.go b/pigo/internal/memory/search_test.go new file mode 100644 index 0000000..d7d7e81 --- /dev/null +++ b/pigo/internal/memory/search_test.go @@ -0,0 +1,152 @@ +package memory + +import "testing" + +// seedSearchCorpus writes a small memory corpus and indexes it. "checkpoint" is +// deliberately a common word (appears in many docs) while "permission" and +// "deadlock" are rare, so BM25 + the relative floor can be exercised. +func seedSearchCorpus(t *testing.T, st *Store, root string) map[string]string { + t.Helper() + paths := map[string]string{} + paths["rare"] = writeFile(t, root, + "permission deadlock encountered during checkpoint save then retry succeeded", + "projects", "proj1", "notes", "rare.md") + paths["c1"] = writeFile(t, root, "checkpoint state alpha", "global", "checkpoint", "c1.md") + paths["c2"] = writeFile(t, root, "checkpoint state beta", "global", "checkpoint", "c2.md") + paths["c3"] = writeFile(t, root, "checkpoint state gamma", "global", "checkpoint", "c3.md") + paths["c4"] = writeFile(t, root, "checkpoint state delta", "global", "checkpoint", "c4.md") + paths["user"] = writeFile(t, root, "unrelated grocery shopping list", "global", "user", "u1.md") + if _, err := st.Reconcile(); err != nil { + t.Fatalf("Reconcile: %v", err) + } + return paths +} + +func resultPaths(rs []SearchResult) map[string]bool { + m := make(map[string]bool, len(rs)) + for _, r := range rs { + m[r.Path] = true + } + return m +} + +func TestSearchMultiWordOrRecall(t *testing.T) { + st, root, _ := openTempWithRoots(t) + p := seedSearchCorpus(t, st, root) + + // OR recall: a query spanning a rare word (in one doc) and the common word + // (in several) should surface docs matching either. Disable the floor so we + // verify raw OR recall independent of trimming. + res, err := st.Search("permission checkpoint", SearchOptions{ScoreFloor: -1}) + if err != nil { + t.Fatalf("Search: %v", err) + } + got := resultPaths(res) + if !got[p["rare"]] { + t.Fatalf("expected rare doc in recall results, got %v", got) + } + if !got[p["c1"]] { + t.Fatalf("expected a checkpoint doc in recall results, got %v", got) + } +} + +func TestSearchScoreFloorDropsCommonWordOnly(t *testing.T) { + st, root, _ := openTempWithRoots(t) + p := seedSearchCorpus(t, st, root) + + // The rare doc matches permission+deadlock+checkpoint; the c* docs match + // only the common "checkpoint". With the default floor the multi-rare doc + // ranks top and the common-word-only docs are trimmed. + res, err := st.Search("permission deadlock checkpoint", SearchOptions{}) + if err != nil { + t.Fatalf("Search: %v", err) + } + if len(res) == 0 { + t.Fatal("expected at least the top result") + } + if res[0].Path != p["rare"] { + t.Fatalf("expected rare doc ranked top, got %q", res[0].Path) + } + // Higher = better after negation: the top score should be positive-most. + for _, r := range res[1:] { + if r.Score > res[0].Score { + t.Fatalf("result %q outscored the top hit", r.Path) + } + } + got := resultPaths(res) + for _, key := range []string{"c1", "c2", "c3", "c4"} { + if got[p[key]] { + t.Fatalf("common-word-only doc %s should have been dropped by floor, got %v", key, got) + } + } +} + +func TestSearchScopeAndTypeFilters(t *testing.T) { + st, root, _ := openTempWithRoots(t) + p := seedSearchCorpus(t, st, root) + + // scope filter: only the projects doc should match under scope=projects. + res, err := st.Search("checkpoint permission", SearchOptions{Scope: "projects", ScoreFloor: -1}) + if err != nil { + t.Fatalf("Search scope: %v", err) + } + got := resultPaths(res) + if !got[p["rare"]] || len(got) != 1 { + t.Fatalf("scope=projects should return only the rare doc, got %v", got) + } + + // type filter: only global/checkpoint docs, none of the projects/user docs. + res, err = st.Search("checkpoint permission", SearchOptions{Type: string(TypeCheckpoint), ScoreFloor: -1}) + if err != nil { + t.Fatalf("Search type: %v", err) + } + got = resultPaths(res) + if got[p["rare"]] || got[p["user"]] { + t.Fatalf("type=checkpoint should exclude non-checkpoint docs, got %v", got) + } + if !got[p["c1"]] { + t.Fatalf("type=checkpoint should include checkpoint docs, got %v", got) + } +} + +func TestSearchEmptyQueryReturnsNil(t *testing.T) { + st, root, _ := openTempWithRoots(t) + seedSearchCorpus(t, st, root) + + for _, q := range []string{"", " ", "!!! ??? ---"} { + res, err := st.Search(q, SearchOptions{}) + if err != nil { + t.Fatalf("Search(%q): unexpected error %v", q, err) + } + if res != nil { + t.Fatalf("Search(%q): expected nil results, got %v", q, res) + } + } +} + +func TestSearchReconcileFirst(t *testing.T) { + st, root, _ := openTempWithRoots(t) + // Write a file but do NOT reconcile manually; ReconcileFirst should index it. + writeFile(t, root, "lazy reconciled permission deadlock content", "global", "notes", "lazy.md") + + res, err := st.Search("permission deadlock", SearchOptions{ReconcileFirst: true}) + if err != nil { + t.Fatalf("Search ReconcileFirst: %v", err) + } + if len(res) == 0 { + t.Fatal("ReconcileFirst should have indexed and matched the lazy doc") + } +} + +func TestSearchLimit(t *testing.T) { + st, root, _ := openTempWithRoots(t) + seedSearchCorpus(t, st, root) + + res, err := st.Search("checkpoint", SearchOptions{Limit: 2, ScoreFloor: -1}) + if err != nil { + t.Fatalf("Search limit: %v", err) + } + if len(res) > 2 { + t.Fatalf("limit=2 should cap results, got %d", len(res)) + } +} diff --git a/pigo/internal/memory/store.go b/pigo/internal/memory/store.go new file mode 100644 index 0000000..3ce8896 --- /dev/null +++ b/pigo/internal/memory/store.go @@ -0,0 +1,137 @@ +// Package memory implements the persistent memory storage layer for pigo. +// +// It is backed by a pure-Go SQLite database (modernc.org/sqlite, no CGO) with +// an FTS5 full-text index over Markdown memory files. This file provides the +// Store type together with database opening and idempotent schema migration. +// Later nodes extend the package with path resolution, reconcile (lazy +// indexing/pruning) and BM25 search. +package memory + +import ( + "database/sql" + "fmt" + "os" + "path/filepath" + + _ "modernc.org/sqlite" // pure-Go SQLite driver, registers the "sqlite" driver +) + +// Store is the handle to the memory database. Fields are unexported; later +// nodes access the underlying *sql.DB via DB() and the configured roots via the +// package-internal fields. +type Store struct { + db *sql.DB + root string // memory root directory (magic layout: global/projects/sessions) + ccBase string // optional Claude Code base dir for the "cc" scope; "" disables +} + +// Open opens (creating if necessary) the SQLite database at dbPath and runs the +// idempotent schema migration. The parent directory of dbPath is created with +// os.MkdirAll before opening. root is the memory root directory and ccBase is +// the optional Claude Code base directory; both are retained for use by later +// nodes and are not required to exist here. +// +// Calling Open twice on the same file is safe: the migration uses +// CREATE ... IF NOT EXISTS throughout. +func Open(dbPath, root, ccBase string) (*Store, error) { + if dbPath == "" { + return nil, fmt.Errorf("memory: empty dbPath") + } + + // modernc.org/sqlite understands the ":memory:" DSN; only create a parent + // directory for real file paths. + if dbPath != ":memory:" { + if dir := filepath.Dir(dbPath); dir != "" && dir != "." { + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("memory: create db dir %q: %w", dir, err) + } + } + } + + db, err := sql.Open("sqlite", dbPath) + if err != nil { + return nil, fmt.Errorf("memory: open db %q: %w", dbPath, err) + } + + // A single global DB with a single connection: writes are serialized, which + // matches the low write frequency and avoids SQLITE_BUSY contention. + db.SetMaxOpenConns(1) + + if err := db.Ping(); err != nil { + db.Close() + return nil, fmt.Errorf("memory: ping db %q: %w", dbPath, err) + } + + if err := migrate(db); err != nil { + db.Close() + return nil, fmt.Errorf("memory: migrate: %w", err) + } + + return &Store{db: db, root: root, ccBase: ccBase}, nil +} + +// migrate applies the idempotent schema DDL. +func migrate(db *sql.DB) error { + if _, err := db.Exec(schemaDDL); err != nil { + return err + } + return nil +} + +// DB returns the underlying *sql.DB for use by later nodes and tests. It may be +// nil if the store was not successfully opened. +func (s *Store) DB() *sql.DB { + if s == nil { + return nil + } + return s.db +} + +// Root returns the memory root directory the store was opened with (the magic +// layout root holding global/projects/sessions). It is the canonical source for +// the memory root used by checkpoint persistence and context rebuild +// (/sessions//checkpoint.md); callers must resolve the memory root +// through this accessor rather than re-deriving it from another store's dir. It +// is "" for a nil store. +func (s *Store) Root() string { + if s == nil { + return "" + } + return s.root +} + +// CountByScope returns the number of indexed memory entries grouped by scope. +// Scopes with no entries are omitted from the map. It reflects the current +// contents of memory_index; callers that want fresh counts should Reconcile +// first. A nil store or nil db yields an empty map. +func (s *Store) CountByScope() (map[Scope]int, error) { + out := make(map[Scope]int) + if s == nil || s.db == nil { + return out, nil + } + rows, err := s.db.Query(`SELECT scope, COUNT(*) FROM memory_index GROUP BY scope`) + if err != nil { + return nil, fmt.Errorf("memory: count by scope: %w", err) + } + defer rows.Close() + for rows.Next() { + var scope string + var n int + if err := rows.Scan(&scope, &n); err != nil { + return nil, fmt.Errorf("memory: scan scope count: %w", err) + } + out[Scope(scope)] = n + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("memory: iterate scope counts: %w", err) + } + return out, nil +} + +// Close closes the underlying database connection. +func (s *Store) Close() error { + if s == nil || s.db == nil { + return nil + } + return s.db.Close() +} diff --git a/pigo/internal/memory/store_test.go b/pigo/internal/memory/store_test.go new file mode 100644 index 0000000..80d5f50 --- /dev/null +++ b/pigo/internal/memory/store_test.go @@ -0,0 +1,175 @@ +package memory + +import ( + "path/filepath" + "testing" +) + +// openTemp opens a Store backed by a temp-file DB and registers cleanup. +func openTemp(t *testing.T) *Store { + t.Helper() + dir := t.TempDir() + dbPath := filepath.Join(dir, "sub", "memory.db") // sub/ must be created by Open + st, err := Open(dbPath, dir, "") + if err != nil { + t.Fatalf("Open: %v", err) + } + t.Cleanup(func() { st.Close() }) + return st +} + +// tableExists reports whether a table or virtual table with the given name +// exists in sqlite_master. +func tableExists(t *testing.T, st *Store, name string) bool { + t.Helper() + var got string + err := st.DB().QueryRow( + `SELECT name FROM sqlite_master WHERE type='table' AND name=?`, name, + ).Scan(&got) + if err != nil { + return false + } + return got == name +} + +// TestFTS5SmokeTest de-risks the pure-Go SQLite FTS5 dependency: it asserts the +// FTS5 virtual table is created without error and is queryable. +func TestFTS5SmokeTest(t *testing.T) { + st := openTemp(t) + + if !tableExists(t, st, "memory_fts") { + t.Fatalf("memory_fts virtual table not created") + } + + // A MATCH query must run without error (proves FTS5 is compiled in). + rows, err := st.DB().Query(`SELECT rowid FROM memory_fts WHERE memory_fts MATCH ?`, "anything") + if err != nil { + t.Fatalf("FTS5 MATCH query failed (FTS5 not available?): %v", err) + } + rows.Close() +} + +// TestSchemaObjectsCreated verifies the content table, indexes and triggers. +func TestSchemaObjectsCreated(t *testing.T) { + st := openTemp(t) + + if !tableExists(t, st, "memory_index") { + t.Fatalf("memory_index table not created") + } + + for _, idx := range []string{"memory_index_scope_idx", "memory_index_type_idx"} { + var n int + if err := st.DB().QueryRow( + `SELECT count(*) FROM sqlite_master WHERE type='index' AND name=?`, idx, + ).Scan(&n); err != nil || n != 1 { + t.Fatalf("index %q missing (n=%d, err=%v)", idx, n, err) + } + } + + for _, trg := range []string{"memory_ai", "memory_ad", "memory_au"} { + var n int + if err := st.DB().QueryRow( + `SELECT count(*) FROM sqlite_master WHERE type='trigger' AND name=?`, trg, + ).Scan(&n); err != nil || n != 1 { + t.Fatalf("trigger %q missing (n=%d, err=%v)", trg, n, err) + } + } +} + +// insertRow inserts a memory_index row and returns its id. +func insertRow(t *testing.T, st *Store, path, body string) int64 { + t.Helper() + res, err := st.DB().Exec( + `INSERT INTO memory_index (path, scope, scope_id, type, body, fingerprint, last_indexed_at) + VALUES (?, 'global', '', 'free', ?, 'fp', 0)`, path, body) + if err != nil { + t.Fatalf("insert: %v", err) + } + id, err := res.LastInsertId() + if err != nil { + t.Fatalf("last insert id: %v", err) + } + return id +} + +// ftsMatchCount returns how many FTS rows match the given single-term query. +func ftsMatchCount(t *testing.T, st *Store, term string) int { + t.Helper() + var n int + if err := st.DB().QueryRow( + `SELECT count(*) FROM memory_fts WHERE memory_fts MATCH ?`, term, + ).Scan(&n); err != nil { + t.Fatalf("fts match count %q: %v", term, err) + } + return n +} + +// TestTriggerSyncInsertDeleteUpdate verifies the AFTER INSERT/DELETE/UPDATE +// triggers keep memory_fts in sync with memory_index. +func TestTriggerSyncInsertDeleteUpdate(t *testing.T) { + st := openTemp(t) + + // INSERT -> searchable. + id := insertRow(t, st, "/mem/a.md", "alpha bravo charlie") + if got := ftsMatchCount(t, st, "bravo"); got != 1 { + t.Fatalf("after insert: MATCH bravo = %d, want 1", got) + } + + // UPDATE -> re-synced (old term gone, new term present). + if _, err := st.DB().Exec(`UPDATE memory_index SET body=? WHERE id=?`, "delta echo", id); err != nil { + t.Fatalf("update: %v", err) + } + if got := ftsMatchCount(t, st, "bravo"); got != 0 { + t.Fatalf("after update: MATCH bravo = %d, want 0", got) + } + if got := ftsMatchCount(t, st, "echo"); got != 1 { + t.Fatalf("after update: MATCH echo = %d, want 1", got) + } + + // DELETE -> removed from index. + if _, err := st.DB().Exec(`DELETE FROM memory_index WHERE id=?`, id); err != nil { + t.Fatalf("delete: %v", err) + } + if got := ftsMatchCount(t, st, "echo"); got != 0 { + t.Fatalf("after delete: MATCH echo = %d, want 0", got) + } +} + +// TestOpenIdempotent verifies that running the migration twice on the same file +// is safe and preserves data. +func TestOpenIdempotent(t *testing.T) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "memory.db") + + st1, err := Open(dbPath, dir, "") + if err != nil { + t.Fatalf("first Open: %v", err) + } + insertRow(t, st1, "/mem/keep.md", "persistent needle") + if err := st1.Close(); err != nil { + t.Fatalf("close first: %v", err) + } + + // Re-open: migration must not error and existing data must survive. + st2, err := Open(dbPath, dir, "") + if err != nil { + t.Fatalf("second Open (migration not idempotent?): %v", err) + } + t.Cleanup(func() { st2.Close() }) + + if got := ftsMatchCount(t, st2, "needle"); got != 1 { + t.Fatalf("after reopen: MATCH needle = %d, want 1", got) + } +} + +// TestInMemoryOpen verifies the ":memory:" DSN works (no parent dir creation). +func TestInMemoryOpen(t *testing.T) { + st, err := Open(":memory:", "", "") + if err != nil { + t.Fatalf("Open in-memory: %v", err) + } + defer st.Close() + if !tableExists(t, st, "memory_fts") { + t.Fatalf("memory_fts not created in in-memory db") + } +} diff --git a/pigo/internal/pihost/embed.go b/pigo/internal/pihost/embed.go new file mode 100644 index 0000000..c06bed5 --- /dev/null +++ b/pigo/internal/pihost/embed.go @@ -0,0 +1,17 @@ +// Package pihost embeds the pi-extension host program (pihost.mjs) into the +// pigo binary. pihost.mjs is a self-contained Node ESM program that loads a pi +// extension using pi's real runtime and re-exposes it over pigo's JSON-RPC +// plugin protocol (see docs/superpowers/specs/2026-07-24-pi-extension-host-design.md). +// +// At install time, internal/pkgmgr.DistributeExtension writes these bytes next +// to a pi extension's payload (plugins/.pkg/.pihost.mjs) and points a +// node-host launcher at them. Embedding keeps the host in lockstep with the +// pigo binary — there is no separate file to ship or version-skew. +package pihost + +import _ "embed" + +// Script is the embedded pihost.mjs source. Callers write it verbatim to disk. +// +//go:embed pihost.mjs +var Script []byte diff --git a/pigo/internal/pihost/host_e2e_test.go b/pigo/internal/pihost/host_e2e_test.go new file mode 100644 index 0000000..27f3ae4 --- /dev/null +++ b/pigo/internal/pihost/host_e2e_test.go @@ -0,0 +1,156 @@ +// End-to-end load test for the embedded pi-extension host (#266). +// +// This exercises the REAL host (pihost.mjs) against the REAL pi SDK +// (@earendil-works/pi-coding-agent) via internal/plugin.Load, driving the +// actual newline-delimited JSON-RPC 2.0 handshake over a subprocess. A tiny +// fixture pi extension is written to a temp dir; the test loads it through the +// host, asserts the fixture's registered command shows up in the manifest, and +// asserts commands/call returns the prompt the fixture emits via +// pi.sendUserMessage (which the host captures into CommandCallResult.Prompt). +// +// It is guarded: it skips cleanly when `node` is not on PATH or when the pi SDK +// cannot be resolved, so CI without a Node/SDK toolchain still passes. +package pihost_test + +import ( + "context" + "encoding/json" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/smallnest/pigo/internal/pihost" + "github.com/smallnest/pigo/internal/plugin" +) + +// piSDKProbe mirrors pihost.mjs's own SDK resolution (loadSdk): try a direct +// import of the package, then fall back to importing the package entry under +// the global npm root. Reporting availability the same way the host resolves +// it keeps the gate honest — the test skips only when the host itself could +// not load the SDK. +const piSDKProbe = ` +const PKG = "@earendil-works/pi-coding-agent"; +import("node:child_process").then(async ({ execFileSync }) => { + try { await import(PKG); process.exit(0); } catch {} + try { + const root = execFileSync("npm", ["root", "-g"], { encoding: "utf8" }).trim(); + const { pathToFileURL } = await import("node:url"); + const path = await import("node:path"); + for (const entry of [ + path.join(root, PKG, "dist", "index.js"), + path.join(root, PKG, "index.js"), + ]) { + try { await import(pathToFileURL(entry).href); process.exit(0); } catch {} + } + } catch {} + process.exit(1); +}).catch(() => process.exit(1)); +` + +// piHostAvailable reports whether the E2E prerequisites are satisfied, and a +// human-readable reason to log when they are not. It requires `node` on PATH +// and a resolvable pi SDK; the SDK probe is bounded so a hung npm cannot stall +// the suite. +func piHostAvailable() (ok bool, reason string) { + if _, err := exec.LookPath("node"); err != nil { + return false, "node is not on PATH" + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, "node", "--input-type=module", "-e", piSDKProbe) + if err := cmd.Run(); err != nil { + return false, "pi SDK (@earendil-works/pi-coding-agent) is not resolvable: " + err.Error() + } + return true, "" +} + +// writeFixtureExtension writes a minimal pi extension package into dir: a +// package.json declaring index.js as its sole pi extension, and an index.js +// that registers a single "e2e" command whose handler emits a known string via +// pi.sendUserMessage. The host captures that message into the command result's +// Prompt. +func writeFixtureExtension(t *testing.T, dir string) { + t.Helper() + + pkgJSON := `{ + "name": "pi-e2e-fixture", + "version": "0.0.0", + "type": "module", + "pi": { "extensions": ["index.js"] } +} +` + indexJS := `export default (pi) => { + pi.registerCommand("e2e", { + description: "e2e probe", + handler: async (args, ctx) => { + pi.sendUserMessage("E2E_PROMPT_OK"); + }, + }); +}; +` + if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(pkgJSON), 0o644); err != nil { + t.Fatalf("write package.json: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "index.js"), []byte(indexJS), 0o644); err != nil { + t.Fatalf("write index.js: %v", err) + } +} + +func TestPiHostExtensionE2E(t *testing.T) { + if ok, reason := piHostAvailable(); !ok { + t.Skip("skipping pi-host E2E: " + reason) + } + + // Exercise the EMBEDDED host bytes: write pihost.Script to a temp .mjs so + // the test drives exactly what pigo ships, not a stray on-disk copy. + tmp := t.TempDir() + hostPath := filepath.Join(tmp, "pihost.mjs") + if err := os.WriteFile(hostPath, pihost.Script, 0o644); err != nil { + t.Fatalf("write embedded pihost.mjs: %v", err) + } + + // The fixture extension lives in its own package dir so the host's pkgDir + // filter (which keeps only extensions under the target dir) selects it. + pkgDir := filepath.Join(tmp, "fixture") + if err := os.MkdirAll(pkgDir, 0o755); err != nil { + t.Fatalf("mkdir fixture: %v", err) + } + writeFixtureExtension(t, pkgDir) + + // Load the extension through the real host. Args mirror the host's argv + // contract: `node pihost.mjs `. plugin.Load performs the initialize + // handshake and decodes the manifest. + p, err := plugin.Load("node", []string{hostPath, pkgDir}, os.Stderr) + if err != nil { + t.Fatalf("plugin.Load(pihost): %v", err) + } + t.Cleanup(func() { _ = p.Close() }) + + // initialize must surface the fixture's registered command. + if got := p.Manifest.Name; got != "pi-e2e-fixture" { + t.Errorf("manifest name = %q, want %q", got, "pi-e2e-fixture") + } + var found bool + for _, c := range p.Manifest.Commands { + if c.Name == "e2e" { + found = true + break + } + } + if !found { + t.Fatalf("manifest commands %+v missing %q", p.Manifest.Commands, "e2e") + } + + // commands/call must run the handler and return the captured prompt. + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + res, err := p.CallCommand(ctx, "e2e", json.RawMessage(`""`)) + if err != nil { + t.Fatalf("CallCommand(e2e): %v", err) + } + if res.Prompt != "E2E_PROMPT_OK" { + t.Fatalf("command prompt = %q (notifications %+v), want %q", res.Prompt, res.Notifications, "E2E_PROMPT_OK") + } +} diff --git a/pigo/internal/pihost/pihost.mjs b/pigo/internal/pihost/pihost.mjs new file mode 100644 index 0000000..601d3fd --- /dev/null +++ b/pigo/internal/pihost/pihost.mjs @@ -0,0 +1,641 @@ +#!/usr/bin/env node +// pihost.mjs — pigo's embedded pi-extension host. +// +// pigo ships this single ESM file (embedded via go:embed, see #264) and writes +// it to disk next to an installed pi extension. pigo launches it as an ordinary +// plugin subprocess and speaks line-delimited JSON-RPC 2.0 over its stdio, +// exactly as it speaks to any other plugin (see internal/plugin, internal/jsonrpc). +// +// The host loads a pi extension using pi's real runtime +// (@earendil-works/pi-coding-agent) and re-exposes its registered tools and +// commands over pigo's plugin protocol: +// +// initialize -> Manifest {name, version, tools[], commands[]} +// tools/call {name, arguments} -> {content, isError} +// commands/call {name, arguments} +// -> {prompt, notifications:[{message,type}]} +// event {type, data} (notify) -> best-effort runner.emit +// shutdown (notify) -> graceful exit +// +// Design invariant: pi actions that pigo does not drive (session/model mutation, +// interactive UI, provider registration, widgets, ...) are inert no-ops that +// NEVER throw. Loading and basic tool/command use must not crash the host. +// +// Argv contract: node pihost.mjs [extra args...] +// Launch cwd: the session cwd (pigo launches the host there). +// +// See docs/superpowers/specs/2026-07-24-pi-extension-host-design.md (§1-§3). + +import { spawn } from "node:child_process"; +import { createInterface } from "node:readline"; +import * as path from "node:path"; +import { pathToFileURL } from "node:url"; + +const PKG = "@earendil-works/pi-coding-agent"; + +// --------------------------------------------------------------------------- +// Diagnostics +// --------------------------------------------------------------------------- + +/** Write one diagnostic line to stderr (pigo pipes plugin stderr through). */ +function diag(msg) { + try { + process.stderr.write(`pihost: ${msg}\n`); + } catch { + // stderr unavailable — nothing more we can do. + } +} + +/** Exit non-zero after a clear diagnostic, before any initialize is answered. */ +function fatal(msg) { + diag(msg); + process.exit(1); +} + +// --------------------------------------------------------------------------- +// SDK resolution (§3) +// +// Import the pi runtime from the public entry without assuming the host file is +// inside a node_modules tree: +// 1. import(PKG) directly (works via NODE_PATH or a local node_modules). +// 2. Fall back to the global npm root (`npm root -g`), then import the +// absolute path to the package entry. +// On failure: clear stderr diagnostic + non-zero exit, BEFORE answering +// initialize, so plugin.Load logs and skips this plugin (fault isolation). +// --------------------------------------------------------------------------- + +/** Run `npm root -g` and return the trimmed path, or "" on any failure. */ +function npmRootGlobal() { + return new Promise((resolve) => { + let out = ""; + let done = false; + const finish = (v) => { + if (!done) { + done = true; + resolve(v); + } + }; + try { + const child = spawn("npm", ["root", "-g"], { stdio: ["ignore", "pipe", "ignore"] }); + child.stdout.on("data", (d) => { + out += d.toString(); + }); + child.on("error", () => finish("")); + child.on("close", (code) => finish(code === 0 ? out.trim() : "")); + // Bound the probe so a hung npm cannot stall startup. + setTimeout(() => { + try { + child.kill(); + } catch { + // ignore + } + finish(""); + }, 5000); + } catch { + finish(""); + } + }); +} + +/** Resolve and import the pi SDK. Returns the module namespace. */ +async function loadSdk() { + // 1. Direct import (NODE_PATH / local node_modules). + try { + return await import(PKG); + } catch (errDirect) { + // 2. Global npm root. + const candidates = []; + const envRoots = (process.env.NODE_PATH || "") + .split(path.delimiter) + .filter(Boolean); + const globalRoot = await npmRootGlobal(); + if (globalRoot) envRoots.unshift(globalRoot); + for (const root of envRoots) { + candidates.push(path.join(root, PKG, "dist", "index.js")); + candidates.push(path.join(root, PKG, "index.js")); + } + for (const entry of candidates) { + try { + return await import(pathToFileURL(entry).href); + } catch { + // try next candidate + } + } + const detail = errDirect instanceof Error ? errDirect.message : String(errDirect); + fatal( + `cannot resolve ${PKG}. Tried direct import and global npm root ` + + `(${globalRoot || "unavailable"}). Is the pi SDK installed? (${detail})`, + ); + } +} + +// --------------------------------------------------------------------------- +// TypeBox parameters -> JSON Schema +// +// A pi tool's `parameters` is a TypeBox TSchema, which is already a plain +// JSON-Schema-shaped object annotated with TypeBox symbol keys ([Kind], etc.). +// JSON.stringify drops symbol-keyed properties, so a round-trip yields a clean +// JSON Schema object. Degrade to a permissive object schema when absent or +// unserializable so tool registration on the pigo side never fails. +// --------------------------------------------------------------------------- +function toJsonSchema(parameters) { + if (parameters && typeof parameters === "object") { + try { + const cleaned = JSON.parse(JSON.stringify(parameters)); + if (cleaned && typeof cleaned === "object") { + if (!cleaned.type) cleaned.type = "object"; + return cleaned; + } + } catch { + // fall through to permissive default + } + } + return { type: "object" }; +} + +// --------------------------------------------------------------------------- +// Content mapping: pi AgentToolResult.content -> plain text +// --------------------------------------------------------------------------- +function contentToText(content) { + if (typeof content === "string") return content; + if (!Array.isArray(content)) return ""; + const parts = []; + for (const c of content) { + if (typeof c === "string") { + parts.push(c); + } else if (c && typeof c === "object") { + if (typeof c.text === "string") { + parts.push(c.text); + } else if (c.type === "image") { + parts.push("[image]"); + } + } + } + return parts.join(""); +} + +// --------------------------------------------------------------------------- +// Command-call capture buffers (§2) +// +// While a commands/call handler runs, pi.sendUserMessage and pi.ui.notify are +// captured into the *current* capture. Outside a command call, ui.notify is +// forwarded to stderr and sendUserMessage is dropped (there is no turn to feed). +// --------------------------------------------------------------------------- +let currentCapture = null; // { prompts: string[], notifications: [{message,type}] } + +function captureUserMessage(content) { + if (!currentCapture) return; + currentCapture.prompts.push(contentToText(content)); +} + +function captureNotify(message, type) { + const msg = typeof message === "string" ? message : String(message ?? ""); + const t = typeof type === "string" && type ? type : "info"; + if (currentCapture) { + currentCapture.notifications.push({ message: msg, type: t }); + } else { + diag(`notify[${t}]: ${msg}`); + } +} + +// --------------------------------------------------------------------------- +// Action bridging (§2) +// +// pigo drives the agent loop, not pi, so these actions are stubs or capture +// buffers. Read-only context values are sensible constants. Mutation actions +// and interactive UI are inert. NOTHING throws. +// --------------------------------------------------------------------------- + +/** UI context: notify captures/forwards; interactive prompts resolve inert. */ +function makeUIContext() { + return { + select: async () => undefined, + confirm: async () => false, + input: async () => undefined, + notify: (message, type) => captureNotify(message, type), + onTerminalInput: () => () => {}, + setStatus: () => {}, + setWorkingMessage: () => {}, + setWorkingVisible: () => {}, + setWorkingIndicator: () => {}, + setHiddenThinkingLabel: () => {}, + setWidget: () => {}, + setFooter: () => {}, + setHeader: () => {}, + setTitle: () => {}, + custom: async () => undefined, + pasteToEditor: () => {}, + setEditorText: () => {}, + getEditorText: () => "", + editor: async () => undefined, + addAutocompleteProvider: () => {}, + setEditorComponent: () => {}, + getEditorComponent: () => undefined, + theme: undefined, + getAllThemes: () => [], + getTheme: () => undefined, + setTheme: () => ({ success: false }), + getToolsExpanded: () => false, + setToolsExpanded: () => {}, + }; +} + +/** ExtensionActions: pi.* action methods. */ +function makeActions() { + return { + sendMessage: () => {}, + sendUserMessage: (content) => captureUserMessage(content), + appendEntry: () => {}, + setSessionName: () => {}, + getSessionName: () => undefined, + setLabel: () => {}, + getActiveTools: () => [], + getAllTools: () => [], + setActiveTools: () => {}, + refreshTools: () => {}, + getCommands: () => [], + setModel: async () => false, + getThinkingLevel: () => "off", + setThinkingLevel: () => {}, + }; +} + +/** ExtensionContextActions: ctx.* in event/tool handlers. Read-only + inert. */ +function makeContextActions() { + return { + getModel: () => undefined, + isIdle: () => true, + isProjectTrusted: () => true, + getSignal: () => undefined, + abort: () => {}, + hasPendingMessages: () => false, + shutdown: () => {}, + getContextUsage: () => undefined, + compact: () => {}, + getSystemPrompt: () => "", + getSystemPromptOptions: () => ({ cwd: process.cwd() }), + }; +} + +// bindCommandContext(undefined) installs safe no-op stubs for the mutation +// handlers (newSession/fork/navigateTree/switchSession/reload/waitForIdle), +// which is exactly the inert behavior we want — see runner.js bindCommandContext. + +// --------------------------------------------------------------------------- +// JSON-RPC 2.0 framing (matches internal/jsonrpc) +// +// Read newline-delimited JSON from stdin; write one JSON object + "\n" per +// response to stdout. A request has an id; a notification omits it. +// --------------------------------------------------------------------------- +function writeMessage(obj) { + try { + process.stdout.write(JSON.stringify(obj) + "\n"); + } catch (err) { + diag(`failed to write response: ${err instanceof Error ? err.message : String(err)}`); + } +} + +function writeResult(id, result) { + writeMessage({ jsonrpc: "2.0", id, result }); +} + +function writeError(id, code, message) { + writeMessage({ jsonrpc: "2.0", id, error: { code, message } }); +} + +// --------------------------------------------------------------------------- +// Host +// --------------------------------------------------------------------------- +async function main() { + const pkgDir = process.argv[2]; + if (!pkgDir) { + fatal("usage: node pihost.mjs [extra args...]"); + } + const absPkgDir = path.resolve(pkgDir); + const cwd = process.cwd(); + + // --- Resolve the SDK (exits non-zero before initialize on failure). --- + const sdk = await loadSdk(); + const { + discoverAndLoadExtensions, + ExtensionRunner, + createEventBus, + } = sdk; + if ( + typeof discoverAndLoadExtensions !== "function" || + typeof ExtensionRunner !== "function" + ) { + fatal( + `${PKG} loaded but its public API is missing expected exports ` + + `(discoverAndLoadExtensions/ExtensionRunner). SDK version mismatch?`, + ); + } + + // --- Load the extension(s) from the package dir via the public loader. --- + // discoverAndLoadExtensions reads the package's pi.extensions and loads the + // declared entrypoints. Pass an event bus so the runtime is fully wired. + const eventBus = typeof createEventBus === "function" ? createEventBus() : undefined; + let loadResult; + try { + loadResult = await discoverAndLoadExtensions([absPkgDir], cwd, undefined, eventBus); + } catch (err) { + fatal( + `failed to load extension from ${absPkgDir}: ` + + `${err instanceof Error ? err.message : String(err)}`, + ); + } + if (loadResult.errors && loadResult.errors.length > 0) { + for (const e of loadResult.errors) { + diag(`extension load error (${e.path}): ${e.error}`); + } + } + + // discoverAndLoadExtensions ALSO scans standard locations (cwd/.pi/extensions + // and the global agent dir), which would merge unrelated extensions into this + // package's manifest. This host re-exposes exactly one installed package, so + // keep only the extensions that live under the target package directory. + const pkgPrefix = absPkgDir + path.sep; + const extensions = (loadResult.extensions || []).filter((ext) => { + const rp = ext && typeof ext.resolvedPath === "string" ? path.resolve(ext.resolvedPath) : ""; + return rp === absPkgDir || rp.startsWith(pkgPrefix); + }); + if (extensions.length === 0) { + fatal(`no pi extensions loaded from ${absPkgDir}`); + } + + // --- Build the runner and bind pigo-appropriate action bridges. --- + // sessionManager/modelRegistry are only touched by mutation paths pigo never + // drives; the runner tolerates minimal stand-ins for tool/command use. + let runner; + try { + runner = new ExtensionRunner( + extensions, + loadResult.runtime, + cwd, + /* sessionManager */ undefined, + /* modelRegistry */ { registerProvider: () => {}, unregisterProvider: () => {} }, + ); + runner.bindCore(makeActions(), makeContextActions(), { + registerProvider: () => {}, + unregisterProvider: () => {}, + }); + runner.bindCommandContext(undefined); + runner.setUIContext(makeUIContext(), "print"); + // Swallow extension-side errors instead of letting them surface as crashes. + if (typeof runner.onError === "function") { + runner.onError((e) => diag(`extension error [${e.event}] ${e.extensionPath}: ${e.error}`)); + } + } catch (err) { + fatal( + `failed to initialize extension runner: ` + + `${err instanceof Error ? err.message : String(err)}`, + ); + } + + // --- Derive the manifest name from the package.json (fallback: dir name). --- + let pkgName = path.basename(absPkgDir); + let pkgVersion = ""; + try { + const { readFileSync } = await import("node:fs"); + const pkg = JSON.parse(readFileSync(path.join(absPkgDir, "package.json"), "utf-8")); + if (pkg && typeof pkg.name === "string" && pkg.name) pkgName = pkg.name; + if (pkg && typeof pkg.version === "string") pkgVersion = pkg.version; + } catch { + // no package.json / unreadable — keep the directory-name fallback. + } + + buildManifestAndServe(runner, pkgName, pkgVersion); +} + +/** Collect the manifest, then serve JSON-RPC over stdio. */ +function buildManifestAndServe(runner, pkgName, pkgVersion) { + // Tools: {name, description, schema(JSON Schema from TypeBox parameters)}. + const registeredTools = safeCall(() => runner.getAllRegisteredTools(), []); + const tools = []; + const toolDefsByName = new Map(); + for (const rt of registeredTools) { + const def = rt && rt.definition ? rt.definition : rt; + if (!def || typeof def.name !== "string") continue; + toolDefsByName.set(def.name, def); + tools.push({ + name: def.name, + description: typeof def.description === "string" ? def.description : "", + schema: toJsonSchema(def.parameters), + }); + } + + // Commands: {name, description, prompt:""}. The prompt is produced at call + // time (from captured sendUserMessage), not declared here. + const registeredCommands = safeCall(() => runner.getRegisteredCommands(), []); + const commands = []; + const commandsByName = new Map(); + for (const rc of registeredCommands) { + // resolveRegisteredCommands returns objects with an invocationName; prefer + // it (it disambiguates duplicate names) and fall back to name. + const invName = rc && (rc.invocationName || rc.name); + if (!invName || typeof rc.handler !== "function") continue; + if (commandsByName.has(invName)) continue; // first registration wins + commandsByName.set(invName, rc); + commands.push({ + name: invName, + description: typeof rc.description === "string" ? rc.description : "", + prompt: "", + }); + } + + const manifest = { name: pkgName, version: pkgVersion, tools, commands }; + + serve(runner, manifest, toolDefsByName, commandsByName); +} + +/** Call fn, returning fallback (and logging) on any throw. */ +function safeCall(fn, fallback) { + try { + return fn(); + } catch (err) { + diag(`recovered from error: ${err instanceof Error ? err.message : String(err)}`); + return fallback; + } +} + +/** Serve JSON-RPC 2.0 requests over stdin/stdout until EOF or shutdown. */ +function serve(runner, manifest, toolDefsByName, commandsByName) { + const rl = createInterface({ input: process.stdin, crlfDelay: Infinity }); + + rl.on("line", (line) => { + const trimmed = line.trim(); + if (!trimmed) return; + let msg; + try { + msg = JSON.parse(trimmed); + } catch { + // Unparseable line: cannot correlate an id, so ignore it (never throw). + diag("ignoring unparseable input line"); + return; + } + // Handle asynchronously; a rejected handler must never crash the host. + handleMessage(msg, runner, manifest, toolDefsByName, commandsByName).catch((err) => { + const id = msg && msg.id !== undefined ? msg.id : null; + diag(`handler error: ${err instanceof Error ? err.message : String(err)}`); + if (id !== null && id !== undefined) { + writeError(id, -32603, "internal error"); + } + }); + }); + + rl.on("close", () => { + // stdin EOF: pigo closed the pipe. Exit cleanly. + process.exit(0); + }); +} + +/** Dispatch one decoded JSON-RPC message. */ +async function handleMessage(msg, runner, manifest, toolDefsByName, commandsByName) { + const { id, method, params } = msg || {}; + const isNotification = id === undefined || id === null; + + switch (method) { + case "initialize": { + if (!isNotification) writeResult(id, manifest); + return; + } + + case "tools/call": { + const result = await callTool(runner, toolDefsByName, params || {}); + if (!isNotification) writeResult(id, result); + return; + } + + case "commands/call": { + const result = await callCommand(runner, commandsByName, params || {}); + if (!isNotification) writeResult(id, result); + return; + } + + case "event": { + // Best-effort lifecycle event delivery; one-way, never throws. + await deliverEvent(runner, params || {}); + return; + } + + case "shutdown": { + // Graceful shutdown. Resolve then exit. + safeCall(() => runner.shutdown && runner.shutdown(), undefined); + process.exit(0); + return; + } + + default: { + // Unsupported method. Reply with an error for requests; ignore for + // notifications. Never throws. + if (!isNotification) { + writeError(id, -32601, `method not found: ${String(method)}`); + } else { + diag(`ignoring unsupported notification: ${String(method)}`); + } + return; + } + } +} + +/** + * Run a pi tool's execute() and map AgentToolResult -> {content, isError}. + * A thrown error, an isError result, or a missing tool -> {isError:true}. + */ +async function callTool(runner, toolDefsByName, params) { + const name = params && params.name; + const args = params && params.arguments !== undefined ? params.arguments : {}; + const def = name ? toolDefsByName.get(name) : undefined; + if (!def || typeof def.execute !== "function") { + return { content: `unknown tool: ${String(name)}`, isError: true }; + } + + const ctx = safeCall(() => runner.createContext(), undefined); + const toolCallId = `pihost-${Date.now()}`; + try { + const result = await def.execute(toolCallId, args, undefined, undefined, ctx); + const content = contentToText(result && result.content); + // A pi tool signals failure by throwing; some also set details.isError. + const isError = Boolean( + result && (result.isError === true || (result.details && result.details.isError === true)), + ); + return { content, isError }; + } catch (err) { + return { + content: `${name}: ${err instanceof Error ? err.message : String(err)}`, + isError: true, + }; + } +} + +/** + * Run a pi command's handler with a capturing context and return + * {prompt, notifications}. prompt = concatenation of captured user messages + * (empty if none). A thrown handler -> a notification + whatever was captured. + */ +async function callCommand(runner, commandsByName, params) { + const name = params && params.name; + const rc = name ? commandsByName.get(name) : undefined; + if (!rc || typeof rc.handler !== "function") { + return { + prompt: "", + notifications: [{ message: `unknown command: ${String(name)}`, type: "error" }], + }; + } + + // Arguments: the plugin protocol passes free-form args as raw JSON under + // "arguments". pi command handlers expect a string (the text after the slash + // command). Coerce: use a string directly, else JSON-encode non-empty values. + let argStr = ""; + const rawArgs = params ? params.arguments : undefined; + if (typeof rawArgs === "string") { + argStr = rawArgs; + } else if (rawArgs !== undefined && rawArgs !== null) { + try { + argStr = JSON.stringify(rawArgs); + } catch { + argStr = ""; + } + } + + const capture = { prompts: [], notifications: [] }; + const previous = currentCapture; + currentCapture = capture; + const ctx = safeCall(() => runner.createCommandContext(), undefined); + try { + await rc.handler(argStr, ctx); + } catch (err) { + capture.notifications.push({ + message: `${name}: ${err instanceof Error ? err.message : String(err)}`, + type: "error", + }); + } finally { + currentCapture = previous; + } + + return { + prompt: capture.prompts.join(""), + notifications: capture.notifications, + }; +} + +/** Best-effort delivery of a lifecycle event to the runner. Never throws. */ +async function deliverEvent(runner, params) { + const type = params && params.type; + if (!type || typeof runner.emit !== "function") return; + let data; + if (params.data !== undefined) { + // data arrives as raw JSON (already parsed by JSON.parse of the line). + data = params.data; + } + const event = data && typeof data === "object" ? { type, ...data } : { type }; + try { + await runner.emit(event); + } catch (err) { + diag(`event ${type} delivery failed: ${err instanceof Error ? err.message : String(err)}`); + } +} + +main().catch((err) => { + fatal(`fatal: ${err instanceof Error ? err.stack || err.message : String(err)}`); +}); diff --git a/pigo/internal/pkgmgr/classify.go b/pigo/internal/pkgmgr/classify.go new file mode 100644 index 0000000..eb2e704 --- /dev/null +++ b/pigo/internal/pkgmgr/classify.go @@ -0,0 +1,150 @@ +// This file classifies a fetched pi package into its type(s) — extension, +// skill, prompt, or theme (#157). Classification reads the package's +// package.json: pi packages carry a "pi" metadata block declaring what they +// provide, and pigo also falls back to structural signals (a bin entry, a +// SKILL.md, a commands/ dir) so a package that omits explicit metadata but +// clearly is one type is still recognized. +// +// A single package may be several types at once — the npm catalog has combined +// "extensionskill" entries — so Classify returns a set. When nothing matches, +// it returns an error rather than guessing, so `pigo install` fails clearly on +// a package that isn't a pi package. +// +// NOTE on metadata shape: the exact pi metadata field names are taken from the +// pi package conventions (a top-level "pi" object with a "type" string or +// "types" array, and/or per-capability keys). Both the explicit "pi.type(s)" +// form and structural fallbacks are honored so classification is robust to +// packages that under-declare. See docs/issue#0157.html for the assumptions. +package pkgmgr + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "slices" +) + +// packageJSON is the subset of an npm package.json pigo reads for classification +// and versioning. Unknown fields are ignored. +type packageJSON struct { + Name string `json:"name"` + Version string `json:"version"` + // Bin is npm's executable declaration: either a string path or a + // {name: path} object. Its presence signals an extension. + Bin json.RawMessage `json:"bin,omitempty"` + // Pi is the pi-specific metadata block declaring the package's capabilities. + Pi *piMeta `json:"pi,omitempty"` +} + +// piMeta is the "pi" block of a package.json. It supports either a single +// "type" or a "types" list, plus per-capability keys so a package can declare, +// e.g., both an extension and a skill. +// +// The pi ecosystem convention (observed across published packages such as +// pi-simplify, pi-mcp-adapter, pi-spark, pi-ask-user) declares capabilities as +// PLURAL arrays of paths — "extensions", "skills", "prompts", "themes" — each +// listing the files/dirs that provide that capability. We also accept the +// singular forms ("extension", "skill", ...) so a package that declares a single +// capability object is still recognized. Any present value (array or object) +// registers that type; only its presence matters for classification. +type piMeta struct { + Type string `json:"type,omitempty"` + Types []string `json:"types,omitempty"` + Extension json.RawMessage `json:"extension,omitempty"` + Extensions json.RawMessage `json:"extensions,omitempty"` + Skill json.RawMessage `json:"skill,omitempty"` + Skills json.RawMessage `json:"skills,omitempty"` + Prompt json.RawMessage `json:"prompt,omitempty"` + Prompts json.RawMessage `json:"prompts,omitempty"` + Theme json.RawMessage `json:"theme,omitempty"` + Themes json.RawMessage `json:"themes,omitempty"` +} + +// Classify inspects the fetched package directory and returns the set of pi +// package types it provides, along with the package name and version read from +// package.json. It returns an error when package.json is missing/unreadable or +// when no known pi type can be determined. +func Classify(pkgDir string) (name, version string, types []PackageType, err error) { + pjPath := filepath.Join(pkgDir, "package.json") + data, err := os.ReadFile(pjPath) + if err != nil { + return "", "", nil, fmt.Errorf("pkgmgr: read package.json: %w", err) + } + var pj packageJSON + if err := json.Unmarshal(data, &pj); err != nil { + return "", "", nil, fmt.Errorf("pkgmgr: parse package.json: %w", err) + } + + set := map[PackageType]bool{} + + // 1. Explicit pi metadata wins. + if pj.Pi != nil { + for _, t := range append(pj.Pi.Types, pj.Pi.Type) { + if pt, ok := normalizeType(t); ok { + set[pt] = true + } + } + if len(pj.Pi.Extension) > 0 || len(pj.Pi.Extensions) > 0 { + set[TypeExtension] = true + } + if len(pj.Pi.Skill) > 0 || len(pj.Pi.Skills) > 0 { + set[TypeSkill] = true + } + if len(pj.Pi.Prompt) > 0 || len(pj.Pi.Prompts) > 0 { + set[TypePrompt] = true + } + if len(pj.Pi.Theme) > 0 || len(pj.Pi.Themes) > 0 { + set[TypeTheme] = true + } + } + + // 2. Structural fallbacks for packages that under-declare. + if len(pj.Bin) > 0 { + set[TypeExtension] = true + } + if fileExists(filepath.Join(pkgDir, "SKILL.md")) { + set[TypeSkill] = true + } + if dirExists(filepath.Join(pkgDir, "commands")) { + set[TypePrompt] = true + } + + if len(set) == 0 { + return "", "", nil, fmt.Errorf("unrecognized pi package: no known pi metadata") + } + + types = make([]PackageType, 0, len(set)) + for t := range set { + types = append(types, t) + } + slices.Sort(types) + return pj.Name, pj.Version, types, nil +} + +// normalizeType maps a pi metadata type string to a PackageType, reporting +// whether it is recognized. +func normalizeType(s string) (PackageType, bool) { + switch PackageType(s) { + case TypeExtension: + return TypeExtension, true + case TypeSkill: + return TypeSkill, true + case TypePrompt: + return TypePrompt, true + case TypeTheme: + return TypeTheme, true + default: + return "", false + } +} + +func fileExists(path string) bool { + info, err := os.Stat(path) + return err == nil && !info.IsDir() +} + +func dirExists(path string) bool { + info, err := os.Stat(path) + return err == nil && info.IsDir() +} diff --git a/pigo/internal/pkgmgr/classify_test.go b/pigo/internal/pkgmgr/classify_test.go new file mode 100644 index 0000000..7470018 --- /dev/null +++ b/pigo/internal/pkgmgr/classify_test.go @@ -0,0 +1,205 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "reflect" + "testing" +) + +// writePkg writes a package.json (and optional extra files) into a fresh temp +// dir and returns the dir. extraFiles maps a relative path to its contents; a +// path ending in "/" is created as a directory. +func writePkg(t *testing.T, packageJSON string, extraFiles map[string]string) string { + t.Helper() + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "package.json"), []byte(packageJSON), 0o644); err != nil { + t.Fatal(err) + } + for name, body := range extraFiles { + p := filepath.Join(dir, name) + if body == "" && name[len(name)-1] == '/' { + if err := os.MkdirAll(p, 0o755); err != nil { + t.Fatal(err) + } + continue + } + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, []byte(body), 0o644); err != nil { + t.Fatal(err) + } + } + return dir +} + +// TestClassifyExplicitType verifies a single explicit pi.type is recognized. +func TestClassifyExplicitType(t *testing.T) { + dir := writePkg(t, `{"name":"pi-web","version":"1.0.0","pi":{"type":"skill"}}`, nil) + name, version, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + if name != "pi-web" || version != "1.0.0" { + t.Errorf("name/version = %q/%q, want pi-web/1.0.0", name, version) + } + if !reflect.DeepEqual(types, []PackageType{TypeSkill}) { + t.Errorf("types = %v, want [skill]", types) + } +} + +// TestClassifyMultiType verifies a package declaring several types via pi.types +// returns them all, sorted. +func TestClassifyMultiType(t *testing.T) { + dir := writePkg(t, `{"name":"combo","version":"2.0.0","pi":{"types":["extension","skill"]}}`, nil) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + want := []PackageType{TypeExtension, TypeSkill} + if !reflect.DeepEqual(types, want) { + t.Errorf("types = %v, want %v", types, want) + } +} + +// TestClassifyPerCapabilityKeys verifies pi.extension + pi.theme blocks both +// register their types. +func TestClassifyPerCapabilityKeys(t *testing.T) { + dir := writePkg(t, `{"name":"x","version":"0.1.0","pi":{"extension":{"cmd":"x"},"theme":{}}}`, nil) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + want := []PackageType{TypeExtension, TypeTheme} + if !reflect.DeepEqual(types, want) { + t.Errorf("types = %v, want %v", types, want) + } +} + +// TestClassifyPluralCapabilityKeys verifies the pi-ecosystem convention of +// plural path arrays (pi.extensions, pi.skills, ...) registers each type. This +// is the shape published packages actually use (pi-simplify, pi-ask-user, ...). +func TestClassifyPluralCapabilityKeys(t *testing.T) { + dir := writePkg(t, `{"name":"pi-ask-user","version":"1.0.0","pi":{"extensions":["./index.ts"],"skills":["./skills"]}}`, nil) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + want := []PackageType{TypeExtension, TypeSkill} + if !reflect.DeepEqual(types, want) { + t.Errorf("types = %v, want %v", types, want) + } +} + +// TestClassifyPluralExtensionsOnly verifies a pure extension declared via +// pi.extensions (no bin) classifies as an extension — the pi-simplify case. +func TestClassifyPluralExtensionsOnly(t *testing.T) { + dir := writePkg(t, `{"name":"pi-simplify","version":"0.2.3","pi":{"extensions":["dist/index.js"]}}`, nil) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + if !reflect.DeepEqual(types, []PackageType{TypeExtension}) { + t.Errorf("types = %v, want [extension]", types) + } +} + +// TestClassifyStructuralBin verifies a bare package with a bin entry is an +// extension even without pi metadata. +func TestClassifyStructuralBin(t *testing.T) { + dir := writePkg(t, `{"name":"pi-mcp-adapter","version":"1.0.0","bin":{"pi-mcp-adapter":"./index.js"}}`, nil) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + if !reflect.DeepEqual(types, []PackageType{TypeExtension}) { + t.Errorf("types = %v, want [extension]", types) + } +} + +// TestClassifyStructuralSkillMd verifies a SKILL.md file signals a skill. +func TestClassifyStructuralSkillMd(t *testing.T) { + dir := writePkg(t, `{"name":"pi-skill","version":"1.0.0"}`, map[string]string{ + "SKILL.md": "# a skill", + }) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + if !reflect.DeepEqual(types, []PackageType{TypeSkill}) { + t.Errorf("types = %v, want [skill]", types) + } +} + +// TestClassifyStructuralCommandsDir verifies a commands/ dir signals a prompt. +func TestClassifyStructuralCommandsDir(t *testing.T) { + dir := writePkg(t, `{"name":"pi-cmds","version":"1.0.0"}`, map[string]string{ + "commands/hello.md": "hi", + }) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + if !reflect.DeepEqual(types, []PackageType{TypePrompt}) { + t.Errorf("types = %v, want [prompt]", types) + } +} + +// TestClassifyExplicitAndStructural verifies explicit metadata and structural +// signals union together (skill via SKILL.md, extension via pi.type). +func TestClassifyExplicitAndStructural(t *testing.T) { + dir := writePkg(t, `{"name":"both","version":"1.0.0","pi":{"type":"extension"}}`, map[string]string{ + "SKILL.md": "# skill", + }) + _, _, types, err := Classify(dir) + if err != nil { + t.Fatalf("Classify: %v", err) + } + want := []PackageType{TypeExtension, TypeSkill} + if !reflect.DeepEqual(types, want) { + t.Errorf("types = %v, want %v", types, want) + } +} + +// TestClassifyUnknown verifies a plain npm package with no pi signals errors. +func TestClassifyUnknown(t *testing.T) { + dir := writePkg(t, `{"name":"lodash","version":"4.17.21"}`, nil) + _, _, _, err := Classify(dir) + if err == nil { + t.Fatal("Classify of non-pi package = nil error, want error") + } + if !contains(err.Error(), "unrecognized pi package") { + t.Errorf("error = %q, want to mention 'unrecognized pi package'", err) + } +} + +// TestClassifyMissingPackageJSON verifies a missing package.json errors. +func TestClassifyMissingPackageJSON(t *testing.T) { + if _, _, _, err := Classify(t.TempDir()); err == nil { + t.Fatal("Classify with no package.json = nil error, want error") + } +} + +// TestClassifyCorruptPackageJSON verifies malformed JSON errors clearly. +func TestClassifyCorruptPackageJSON(t *testing.T) { + dir := writePkg(t, `{not valid json`, nil) + _, _, _, err := Classify(dir) + if err == nil { + t.Fatal("Classify with corrupt package.json = nil error, want error") + } + if !contains(err.Error(), "parse package.json") { + t.Errorf("error = %q, want to mention 'parse package.json'", err) + } +} + +// TestClassifyUnknownTypeStringIgnored verifies an unrecognized pi.type string +// is ignored (not fatal) but leaves the package unclassified if nothing else +// matches. +func TestClassifyUnknownTypeStringIgnored(t *testing.T) { + dir := writePkg(t, `{"name":"x","version":"1.0.0","pi":{"type":"widget"}}`, nil) + _, _, _, err := Classify(dir) + if err == nil { + t.Fatal("Classify with only unknown pi.type = nil error, want error") + } +} diff --git a/pigo/internal/pkgmgr/distribute.go b/pigo/internal/pkgmgr/distribute.go new file mode 100644 index 0000000..5612f09 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute.go @@ -0,0 +1,260 @@ +// This file distributes a classified pi extension (including MCP adapters) into +// pigo's plugins directory so internal/plugin.Discover picks it up (#158). +// +// plugin.Discover launches every *executable regular file* directly inside +// $PIGO_HOME/plugins, ignoring subdirectories. An npm extension, however, is a +// whole package tree with a "bin" entry pointing at its real entrypoint, and +// that entrypoint usually needs its sibling files present to run. So we cannot +// just drop a single file in. +// +// The layout we lay down reconciles the two: +// +// $PIGO_HOME/plugins/.pkg/ ← full extracted package (a dir; Discover skips it) +// $PIGO_HOME/plugins/ ← executable launcher (a file; Discover runs it) +// +// The launcher is a tiny shell script that execs the package's bin entrypoint, +// forwarding argv. The bin file is made executable and relies on its own +// shebang (matching how npm itself installs bins), so both Node scripts and +// native binaries work. Every file laid down is returned so the lockfile can +// remove exactly what was created on uninstall. +package pkgmgr + +import ( + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "runtime" + "strings" + + "github.com/smallnest/pigo/internal/pihost" +) + +// DistributeExtension copies the extension package at pkgDir into the plugins +// directory and writes a launcher that plugin.Discover will run. It returns the +// absolute paths of every file (and the payload dir) it created, for the +// lockfile. An empty plugins dir (home unavailable) is an error. +func DistributeExtension(pkgDir, name string) ([]string, error) { + if runtime.GOOS == "windows" { + return nil, fmt.Errorf("pkgmgr: extension install is not supported on windows yet") + } + pluginsDir := PluginsDir() + if pluginsDir == "" { + return nil, fmt.Errorf("pkgmgr: cannot resolve plugins dir (PIGO_HOME/home unavailable)") + } + + binRel, err := extensionBin(pkgDir, name) + if err != nil { + return nil, err + } + + if err := os.MkdirAll(pluginsDir, 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: create plugins dir: %w", err) + } + + payloadDir := filepath.Join(pluginsDir, name+".pkg") + // A stale payload from a prior install must not shadow the new one. + if err := os.RemoveAll(payloadDir); err != nil { + return nil, fmt.Errorf("pkgmgr: clear old payload %q: %w", payloadDir, err) + } + created, err := copyTree(pkgDir, payloadDir) + if err != nil { + return nil, err + } + + // Make the bin entrypoint executable; it carries its own shebang. + binAbs := filepath.Join(payloadDir, filepath.FromSlash(binRel)) + if err := os.Chmod(binAbs, 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: chmod bin %q: %w", binAbs, err) + } + + launcher := filepath.Join(pluginsDir, name) + + // A pi extension is a JS module loaded by pi's runtime, not a native binary + // or JSON-RPC server, so it cannot be exec'd directly. Instead we drop the + // embedded Node host next to the payload and point the launcher at it. A + // native/JSON-RPC plugin keeps the historical direct-exec launcher. + if isPiExtension(pkgDir, binRel) { + hostAbs := filepath.Join(payloadDir, ".pihost.mjs") + if err := os.WriteFile(hostAbs, pihost.Script, 0o644); err != nil { + return nil, fmt.Errorf("pkgmgr: write pi host %q: %w", hostAbs, err) + } + script := fmt.Sprintf( + "#!/bin/sh\n"+ + "command -v node >/dev/null 2>&1 || { echo \"pigo: node not found on PATH; pi extension %s skipped\" >&2; exit 127; }\n"+ + "exec node %s %s \"$@\"\n", + name, shellQuote(hostAbs), shellQuote(payloadDir), + ) + if err := os.WriteFile(launcher, []byte(script), 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: write launcher %q: %w", launcher, err) + } + created = append(created, payloadDir, hostAbs, launcher) + return created, nil + } + + script := fmt.Sprintf("#!/bin/sh\nexec %s \"$@\"\n", shellQuote(binAbs)) + if err := os.WriteFile(launcher, []byte(script), 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: write launcher %q: %w", launcher, err) + } + + created = append(created, payloadDir, launcher) + return created, nil +} + +// isPiExtension reports whether the entrypoint should be run under pigo's Node +// host rather than exec'd directly. The pi.extensions declaration is +// authoritative and checked first; otherwise the resolved bin extension +// (.js/.mjs/.cjs) identifies a JS module that needs the host. Native binaries +// and self-hosted JSON-RPC servers match neither and keep the direct launcher. +func isPiExtension(pkgDir, binRel string) bool { + if data, err := os.ReadFile(filepath.Join(pkgDir, "package.json")); err == nil { + var pj struct { + Pi struct { + Extensions []string `json:"extensions"` + } `json:"pi"` + } + if json.Unmarshal(data, &pj) == nil { + for _, p := range pj.Pi.Extensions { + if p != "" { + return true + } + } + } + } + switch strings.ToLower(filepath.Ext(binRel)) { + case ".js", ".mjs", ".cjs": + return true + } + return false +} + +// extensionBin resolves the package's entrypoint (relative to the package root) +// from package.json. It prefers npm's "bin" field (a string, or a {command: +// path} object keyed by the package name), matching how npm installs bins. When +// there is no "bin" — the common case for pi extensions, which declare their +// entrypoint in the pi metadata rather than as an npm bin — it falls back to the +// first path listed in "pi.extensions", then to "main". +func extensionBin(pkgDir, name string) (string, error) { + data, err := os.ReadFile(filepath.Join(pkgDir, "package.json")) + if err != nil { + return "", fmt.Errorf("pkgmgr: read package.json: %w", err) + } + var pj struct { + Bin json.RawMessage `json:"bin"` + Main string `json:"main"` + Pi struct { + Extensions []string `json:"extensions"` + } `json:"pi"` + } + if err := json.Unmarshal(data, &pj); err != nil { + return "", fmt.Errorf("pkgmgr: parse package.json: %w", err) + } + + // 1. npm "bin": string form. + var s string + if err := json.Unmarshal(pj.Bin, &s); err == nil && s != "" { + return s, nil + } + // npm "bin": object form {command: path}. + var m map[string]string + if err := json.Unmarshal(pj.Bin, &m); err == nil && len(m) > 0 { + if p, ok := m[name]; ok && p != "" { + return p, nil + } + // npm-scoped name: bin key is often the unscoped base name. + if p, ok := m[filepath.Base(name)]; ok && p != "" { + return p, nil + } + for _, p := range m { + if p != "" { + return p, nil + } + } + } + + // 2. pi metadata entrypoint: pi.extensions is the pi-ecosystem convention. + for _, p := range pj.Pi.Extensions { + if p != "" { + return p, nil + } + } + + // 3. npm "main" as a last resort. + if pj.Main != "" { + return pj.Main, nil + } + + return "", fmt.Errorf("pkgmgr: extension %q has no bin, pi.extensions, or main entrypoint in package.json", name) +} + +// copyTree recursively copies src into dst (created), preserving file modes and +// relative structure. It returns the absolute paths of every regular file it +// wrote (not directories), so callers can record precisely what was laid down. +func copyTree(src, dst string) ([]string, error) { + var files []string + err := filepath.Walk(src, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + rel, err := filepath.Rel(src, path) + if err != nil { + return err + } + target := filepath.Join(dst, rel) + switch { + case info.IsDir(): + return os.MkdirAll(target, 0o755) + case info.Mode().IsRegular(): + if err := copyFile(path, target, info.Mode()); err != nil { + return err + } + files = append(files, target) + return nil + default: + // Skip symlinks/devices — npm packages are files + dirs. + return nil + } + }) + if err != nil { + return nil, fmt.Errorf("pkgmgr: copy package tree: %w", err) + } + return files, nil +} + +// copyFile copies a single regular file from src to dst with the given mode. +func copyFile(src, dst string, mode os.FileMode) error { + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + return err + } + in, err := os.Open(src) + if err != nil { + return err + } + defer in.Close() + out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, mode) + if err != nil { + return err + } + if _, err := io.Copy(out, in); err != nil { + out.Close() + return err + } + return out.Close() +} + +// shellQuote wraps s in single quotes for safe embedding in a /bin/sh script, +// escaping any embedded single quotes. +func shellQuote(s string) string { + quoted := make([]byte, 0, len(s)+2) + quoted = append(quoted, '\'') + for i := 0; i < len(s); i++ { + if s[i] == '\'' { + quoted = append(quoted, '\'', '\\', '\'', '\'') + continue + } + quoted = append(quoted, s[i]) + } + quoted = append(quoted, '\'') + return string(quoted) +} diff --git a/pigo/internal/pkgmgr/distribute_prompt.go b/pigo/internal/pkgmgr/distribute_prompt.go new file mode 100644 index 0000000..2f26931 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_prompt.go @@ -0,0 +1,85 @@ +// This file distributes a classified pi prompt/command package into pigo's +// prompts directory so runtime.LoadUserCommandsDir picks it up (#160, #342). +// +// pigo loads declarative slash commands from $PIGO_HOME/prompts/*.md (and the +// legacy $PIGO_HOME/commands/*.md) non-recursively: each markdown file defines +// a "/name" command, named after the file, whose body is the prompt template. +// A pi prompt package ships one or more such templates, conventionally under a +// "prompts/" subdirectory (the pi convention); the legacy "commands/" subdir is +// a fallback, and some packages place the .md files at the package root. +// +// Distribution copies those .md files (flattened, since the loader is +// non-recursive) into $PIGO_HOME/prompts/. Every file laid down is returned so +// the lockfile can remove precisely what was installed. +package pkgmgr + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +// DistributePrompt copies the prompt package's command templates at pkgDir into +// the commands directory. It looks first in "/commands" for *.md files +// and falls back to *.md at the package root. It returns the absolute paths of +// every file it created, for the lockfile. An unresolvable commands dir, or a +// package with no command templates, is an error. +func DistributePrompt(pkgDir, name string) ([]string, error) { + promptsDir := PromptsDir() + if promptsDir == "" { + return nil, fmt.Errorf("pkgmgr: cannot resolve prompts dir (PIGO_HOME/home unavailable)") + } + + // Prefer the pi-aligned prompts/ subdir, then the legacy commands/ subdir, + // then root-level *.md as a last resort. + srcDir := "" + for _, sub := range []string{"prompts", "commands"} { + if dirExists(filepath.Join(pkgDir, sub)) { + srcDir = filepath.Join(pkgDir, sub) + break + } + } + if srcDir == "" { + srcDir = pkgDir // fall back to root-level *.md + } + entries, err := os.ReadDir(srcDir) + if err != nil { + return nil, fmt.Errorf("pkgmgr: read prompt source dir: %w", err) + } + + var mds []string + for _, e := range entries { + if e.IsDir() || !strings.EqualFold(filepath.Ext(e.Name()), ".md") { + continue + } + // A root-level README.md is not a command template; skip it when we've + // fallen back to the package root. + if srcDir == pkgDir && strings.EqualFold(e.Name(), "README.md") { + continue + } + mds = append(mds, e.Name()) + } + if len(mds) == 0 { + return nil, fmt.Errorf("pkgmgr: prompt %q has no prompt templates (*.md)", name) + } + + if err := os.MkdirAll(promptsDir, 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: create prompts dir: %w", err) + } + + created := make([]string, 0, len(mds)) + for _, md := range mds { + src := filepath.Join(srcDir, md) + dst := filepath.Join(promptsDir, md) + info, statErr := os.Stat(src) + if statErr != nil { + return created, fmt.Errorf("pkgmgr: stat command %q: %w", src, statErr) + } + if err := copyFile(src, dst, info.Mode()); err != nil { + return created, fmt.Errorf("pkgmgr: copy command %q: %w", md, err) + } + created = append(created, dst) + } + return created, nil +} diff --git a/pigo/internal/pkgmgr/distribute_prompt_test.go b/pigo/internal/pkgmgr/distribute_prompt_test.go new file mode 100644 index 0000000..ecf631c --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_prompt_test.go @@ -0,0 +1,121 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "testing" +) + +// TestDistributePromptPromptsDirPreferred verifies the pi-aligned prompts/ +// subdir is preferred over the legacy commands/ subdir (#342): only the +// prompts/ templates are installed. +func TestDistributePromptPromptsDirPreferred(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-prompts","version":"1.0.0","pi":{"type":"prompt"}}`, map[string]string{ + "prompts/review.md": "---\ndescription: review\n---\nReview $ARGUMENTS", + "commands/legacy.md": "Legacy $ARGUMENTS", + }) + + files, err := DistributePrompt(pkg, "pi-prompts") + if err != nil { + t.Fatalf("DistributePrompt: %v", err) + } + if len(files) != 1 { + t.Fatalf("created %d files, want 1 (prompts/ preferred over commands/): %v", len(files), files) + } + if _, err := os.Stat(filepath.Join(home, "prompts", "review.md")); err != nil { + t.Errorf("review.md not placed in prompts/: %v", err) + } + if _, err := os.Stat(filepath.Join(home, "prompts", "legacy.md")); !os.IsNotExist(err) { + t.Errorf("legacy.md from commands/ should not be installed when prompts/ exists: %v", err) + } +} + +// TestDistributePromptCommandsDirFallback verifies the legacy commands/ subdir +// is used when prompts/ is absent, installing into ~/.pigo/prompts. +func TestDistributePromptCommandsDirFallback(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-prompts","version":"1.0.0","pi":{"type":"prompt"}}`, map[string]string{ + "commands/review.md": "---\ndescription: review\n---\nReview $ARGUMENTS", + "commands/explain.md": "Explain $ARGUMENTS", + "README.md": "readme", + }) + + files, err := DistributePrompt(pkg, "pi-prompts") + if err != nil { + t.Fatalf("DistributePrompt: %v", err) + } + if len(files) != 2 { + t.Errorf("created %d files, want 2: %v", len(files), files) + } + for _, want := range []string{"review.md", "explain.md"} { + if _, err := os.Stat(filepath.Join(home, "prompts", want)); err != nil { + t.Errorf("%s not placed in prompts/: %v", want, err) + } + } + // Root README.md is not copied (commands/ was used, not root). + if _, err := os.Stat(filepath.Join(home, "prompts", "README.md")); !os.IsNotExist(err) { + t.Errorf("root README.md leaked into prompts: %v", err) + } +} + +// TestDistributePromptRootFallback verifies root-level *.md are used when +// neither prompts/ nor commands/ exists, skipping README.md, installing into +// ~/.pigo/prompts. +func TestDistributePromptRootFallback(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-p","version":"1.0.0","pi":{"type":"prompt"}}`, map[string]string{ + "summarize.md": "Summarize $ARGUMENTS", + "README.md": "readme", + }) + + files, err := DistributePrompt(pkg, "pi-p") + if err != nil { + t.Fatalf("DistributePrompt: %v", err) + } + if len(files) != 1 { + t.Fatalf("created %d files, want 1: %v", len(files), files) + } + if _, err := os.Stat(filepath.Join(home, "prompts", "summarize.md")); err != nil { + t.Errorf("summarize.md not placed in prompts/: %v", err) + } + if _, err := os.Stat(filepath.Join(home, "prompts", "README.md")); !os.IsNotExist(err) { + t.Errorf("README.md should be skipped at root fallback: %v", err) + } +} + +// TestDistributePromptNone verifies a package with no prompt templates errors. +func TestDistributePromptNone(t *testing.T) { + t.Setenv("PIGO_HOME", t.TempDir()) + pkg := writePkg(t, `{"name":"pi-empty","version":"1.0.0"}`, nil) + if _, err := DistributePrompt(pkg, "pi-empty"); err == nil { + t.Fatal("DistributePrompt with no templates = nil error, want error") + } +} + +// TestDistributePromptReturnsPromptsPaths verifies the returned paths (recorded +// in the lockfile) are under ~/.pigo/prompts, so uninstall removes precisely +// what was installed. +func TestDistributePromptReturnsPromptsPaths(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + pkg := writePkg(t, `{"name":"pi-p","version":"1.0.0","pi":{"type":"prompt"}}`, map[string]string{ + "prompts/x.md": "X $ARGUMENTS", + }) + files, err := DistributePrompt(pkg, "pi-p") + if err != nil { + t.Fatalf("DistributePrompt: %v", err) + } + wantDir := filepath.Join(home, "prompts") + for _, f := range files { + if filepath.Dir(f) != wantDir { + t.Errorf("installed path %q not under %s", f, wantDir) + } + } +} diff --git a/pigo/internal/pkgmgr/distribute_skill.go b/pigo/internal/pkgmgr/distribute_skill.go new file mode 100644 index 0000000..176dd63 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_skill.go @@ -0,0 +1,48 @@ +// This file distributes a classified pi skill into pigo's skills directory so +// runtime.LoadSkillsDir picks it up (#159). +// +// pigo loads skills from the skills dir (SkillsDir): ~/.agents/skills, or +// PIGO_SKILLS_DIR when set. LoadSkillsDir recognizes the nested layout +// "//SKILL.md" (a directory per skill whose SKILL.md holds the +// YAML frontmatter). An npm skill package is exactly such a bundle — a SKILL.md +// plus its supporting files — so distribution is a straight copy of the package +// tree into "//". +// +// As with extensions, every file laid down is returned so the lockfile can +// remove precisely what was installed on uninstall. +package pkgmgr + +import ( + "fmt" + "os" + "path/filepath" +) + +// DistributeSkill copies the skill package at pkgDir into the skills directory +// under a "/" subdirectory, where runtime.LoadSkillsDir discovers it via +// its SKILL.md. It returns the absolute paths of every file (and the skill dir) +// it created, for the lockfile. An unresolvable skills dir, or a package with no +// SKILL.md, is an error. +func DistributeSkill(pkgDir, name string) ([]string, error) { + skillsDir := SkillsDir() + if skillsDir == "" { + return nil, fmt.Errorf("pkgmgr: cannot resolve skills dir (home unavailable)") + } + if !fileExists(filepath.Join(pkgDir, "SKILL.md")) { + return nil, fmt.Errorf("pkgmgr: skill %q has no SKILL.md", name) + } + + dest := filepath.Join(skillsDir, name) + // A stale skill from a prior install must not linger alongside the new one. + if err := os.RemoveAll(dest); err != nil { + return nil, fmt.Errorf("pkgmgr: clear old skill %q: %w", dest, err) + } + if err := os.MkdirAll(skillsDir, 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: create skills dir: %w", err) + } + files, err := copyTree(pkgDir, dest) + if err != nil { + return nil, err + } + return append(files, dest), nil +} diff --git a/pigo/internal/pkgmgr/distribute_skill_test.go b/pigo/internal/pkgmgr/distribute_skill_test.go new file mode 100644 index 0000000..7b42513 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_skill_test.go @@ -0,0 +1,76 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "testing" +) + +// TestDistributeSkill verifies a skill package copies into skillsDir// +// with its SKILL.md and supporting files, discoverable by LoadSkillsDir's +// nested layout. +func TestDistributeSkill(t *testing.T) { + skills := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", skills) + + pkg := writePkg(t, `{"name":"pi-writing","version":"1.0.0","pi":{"type":"skill"}}`, map[string]string{ + "SKILL.md": "---\nname: writing\ndescription: help writing\n---\nbody", + "references/tips.md": "tips", + }) + + files, err := DistributeSkill(pkg, "pi-writing") + if err != nil { + t.Fatalf("DistributeSkill: %v", err) + } + + // SKILL.md must land at //SKILL.md (nested layout). + skillMd := filepath.Join(skills, "pi-writing", "SKILL.md") + if _, err := os.Stat(skillMd); err != nil { + t.Fatalf("SKILL.md not placed: %v", err) + } + if _, err := os.Stat(filepath.Join(skills, "pi-writing", "references", "tips.md")); err != nil { + t.Errorf("supporting file not copied: %v", err) + } + + var sawSkillMd bool + for _, f := range files { + if f == skillMd { + sawSkillMd = true + } + } + if !sawSkillMd { + t.Errorf("created files %v missing SKILL.md", files) + } +} + +// TestDistributeSkillReinstallReplaces verifies reinstalling clears stale files. +func TestDistributeSkillReinstallReplaces(t *testing.T) { + skills := t.TempDir() + t.Setenv("PIGO_SKILLS_DIR", skills) + + pkg1 := writePkg(t, `{"name":"pi-s","version":"1.0.0"}`, map[string]string{ + "SKILL.md": "---\nname: s\ndescription: d\n---\n", + "old.md": "old", + }) + if _, err := DistributeSkill(pkg1, "pi-s"); err != nil { + t.Fatalf("first install: %v", err) + } + pkg2 := writePkg(t, `{"name":"pi-s","version":"2.0.0"}`, map[string]string{ + "SKILL.md": "---\nname: s\ndescription: d2\n---\n", + }) + if _, err := DistributeSkill(pkg2, "pi-s"); err != nil { + t.Fatalf("reinstall: %v", err) + } + if _, err := os.Stat(filepath.Join(skills, "pi-s", "old.md")); !os.IsNotExist(err) { + t.Errorf("stale skill file survived reinstall: %v", err) + } +} + +// TestDistributeSkillNoSkillMd verifies a package without SKILL.md errors. +func TestDistributeSkillNoSkillMd(t *testing.T) { + t.Setenv("PIGO_SKILLS_DIR", t.TempDir()) + pkg := writePkg(t, `{"name":"pi-noskill","version":"1.0.0"}`, nil) + if _, err := DistributeSkill(pkg, "pi-noskill"); err == nil { + t.Fatal("DistributeSkill without SKILL.md = nil error, want error") + } +} diff --git a/pigo/internal/pkgmgr/distribute_test.go b/pigo/internal/pkgmgr/distribute_test.go new file mode 100644 index 0000000..f7bce47 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_test.go @@ -0,0 +1,264 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "runtime" + "testing" +) + +// TestDistributeExtensionStringBin verifies a package with a string "bin" +// installs a launcher + payload tree and reports created files. +func TestDistributeExtensionStringBin(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-demo","version":"1.0.0","bin":"./cli.js"}`, map[string]string{ + "cli.js": "#!/usr/bin/env node\nconsole.log('hi')\n", + "lib/x.js": "module.exports=1\n", + }) + + files, err := DistributeExtension(pkg, "pi-demo") + if err != nil { + t.Fatalf("DistributeExtension: %v", err) + } + + launcher := filepath.Join(home, "plugins", "pi-demo") + info, err := os.Stat(launcher) + if err != nil { + t.Fatalf("stat launcher: %v", err) + } + if info.Mode()&0o111 == 0 { + t.Errorf("launcher not executable: mode %v", info.Mode()) + } + + // The payload bin must exist and be executable. + binAbs := filepath.Join(home, "plugins", "pi-demo.pkg", "cli.js") + bi, err := os.Stat(binAbs) + if err != nil { + t.Fatalf("stat payload bin: %v", err) + } + if bi.Mode()&0o111 == 0 { + t.Errorf("payload bin not executable: mode %v", bi.Mode()) + } + + // Sibling files copied. + if _, err := os.Stat(filepath.Join(home, "plugins", "pi-demo.pkg", "lib", "x.js")); err != nil { + t.Errorf("sibling file not copied: %v", err) + } + + // created list includes launcher and payload dir. + var sawLauncher bool + for _, f := range files { + if f == launcher { + sawLauncher = true + } + } + if !sawLauncher { + t.Errorf("created files %v missing launcher", files) + } + + // A ".js" bin is a pi extension, so the launcher runs the Node host and a + // .pihost.mjs is dropped beside the payload. + host := filepath.Join(home, "plugins", "pi-demo.pkg", ".pihost.mjs") + if _, err := os.Stat(host); err != nil { + t.Errorf("expected embedded pi host at %q: %v", host, err) + } + script, _ := os.ReadFile(launcher) + if !contains(string(script), "exec node ") || !contains(string(script), host) { + t.Errorf("launcher script = %q, want node host exec of %q", script, host) + } +} + +// TestDistributeExtensionObjectBin verifies the {command: path} bin form, +// preferring the entry keyed by package name. +func TestDistributeExtensionObjectBin(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-adapter","version":"2.0.0","bin":{"pi-adapter":"./main.js","other":"./other.js"}}`, map[string]string{ + "main.js": "#!/usr/bin/env node\n", + "other.js": "#!/usr/bin/env node\n", + }) + + if _, err := DistributeExtension(pkg, "pi-adapter"); err != nil { + t.Fatalf("DistributeExtension: %v", err) + } + + binAbs := filepath.Join(home, "plugins", "pi-adapter.pkg", "main.js") + if bi, err := os.Stat(binAbs); err != nil || bi.Mode()&0o111 == 0 { + t.Errorf("expected main.js executable, err=%v", err) + } +} + +// TestDistributeExtensionReinstallReplaces verifies a second install clears the +// stale payload rather than merging it. +func TestDistributeExtensionReinstallReplaces(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg1 := writePkg(t, `{"name":"pi-x","version":"1.0.0","bin":"./a.js"}`, map[string]string{ + "a.js": "#!/usr/bin/env node\n", + "gone.js": "old\n", + }) + if _, err := DistributeExtension(pkg1, "pi-x"); err != nil { + t.Fatalf("first install: %v", err) + } + + pkg2 := writePkg(t, `{"name":"pi-x","version":"2.0.0","bin":"./a.js"}`, map[string]string{ + "a.js": "#!/usr/bin/env node\n", + }) + if _, err := DistributeExtension(pkg2, "pi-x"); err != nil { + t.Fatalf("reinstall: %v", err) + } + + // The file only present in the first install must be gone. + if _, err := os.Stat(filepath.Join(home, "plugins", "pi-x.pkg", "gone.js")); !os.IsNotExist(err) { + t.Errorf("stale file survived reinstall: %v", err) + } +} + +// TestDistributeExtensionPiExtensionsEntry verifies a pi extension with no npm +// "bin" resolves its entrypoint from pi.extensions (the pi-simplify shape). +func TestDistributeExtensionPiExtensionsEntry(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-simplify","version":"0.2.3","pi":{"extensions":["dist/index.js"]}}`, map[string]string{ + "dist/index.js": "#!/usr/bin/env node\nconsole.log('hi')\n", + }) + + files, err := DistributeExtension(pkg, "pi-simplify") + if err != nil { + t.Fatalf("DistributeExtension: %v", err) + } + + binAbs := filepath.Join(home, "plugins", "pi-simplify.pkg", "dist", "index.js") + bi, err := os.Stat(binAbs) + if err != nil { + t.Fatalf("stat payload bin: %v", err) + } + if bi.Mode()&0o111 == 0 { + t.Errorf("payload bin not executable: mode %v", bi.Mode()) + } + + // A pi.extensions package must run under the Node host, dropping .pihost.mjs + // and pointing the launcher at `node `. + launcher := filepath.Join(home, "plugins", "pi-simplify") + host := filepath.Join(home, "plugins", "pi-simplify.pkg", ".pihost.mjs") + pkgDir := filepath.Join(home, "plugins", "pi-simplify.pkg") + if _, err := os.Stat(host); err != nil { + t.Errorf("expected embedded pi host at %q: %v", host, err) + } + script, _ := os.ReadFile(launcher) + if !contains(string(script), "exec node ") || !contains(string(script), host) || !contains(string(script), pkgDir) { + t.Errorf("launcher script = %q, want node host exec of %q with pkgDir %q", script, host, pkgDir) + } + if !contains(string(script), "node not found") { + t.Errorf("launcher script = %q, missing node-absent guard", script) + } + + // Both the host and the launcher must be recorded for uninstall. + var sawHost, sawLauncher bool + for _, f := range files { + switch f { + case host: + sawHost = true + case launcher: + sawLauncher = true + } + } + if !sawHost || !sawLauncher { + t.Errorf("created files %v missing host or launcher", files) + } +} + +// TestDistributeExtensionBinaryBinDirectExec verifies a native binary bin (no +// pi.extensions, non-JS extension) keeps the historical direct-exec launcher — +// no Node host, no .pihost.mjs. +func TestDistributeExtensionBinaryBinDirectExec(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-native","version":"1.0.0","bin":"./server"}`, map[string]string{ + "server": "#!/usr/bin/env node\n", + }) + + files, err := DistributeExtension(pkg, "pi-native") + if err != nil { + t.Fatalf("DistributeExtension: %v", err) + } + + launcher := filepath.Join(home, "plugins", "pi-native") + binAbs := filepath.Join(home, "plugins", "pi-native.pkg", "server") + script, _ := os.ReadFile(launcher) + if !contains(string(script), binAbs) { + t.Errorf("launcher script = %q, want direct exec of %q", script, binAbs) + } + if contains(string(script), "exec node ") { + t.Errorf("binary bin launcher = %q, should not run the Node host", script) + } + if _, err := os.Stat(filepath.Join(home, "plugins", "pi-native.pkg", ".pihost.mjs")); !os.IsNotExist(err) { + t.Errorf("binary bin should not drop .pihost.mjs, err=%v", err) + } + + // The launcher and payload dir must be recorded; the host must not be. + var sawLauncher bool + for _, f := range files { + if f == launcher { + sawLauncher = true + } + } + if !sawLauncher { + t.Errorf("created files %v missing launcher", files) + } +} + +// TestDistributeExtensionMainEntry verifies the "main" field is used as a +// last-resort entrypoint when neither bin nor pi.extensions is present. +func TestDistributeExtensionMainEntry(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-main","version":"1.0.0","main":"./dist/index.js"}`, map[string]string{ + "dist/index.js": "#!/usr/bin/env node\n", + }) + + if _, err := DistributeExtension(pkg, "pi-main"); err != nil { + t.Fatalf("DistributeExtension: %v", err) + } + binAbs := filepath.Join(home, "plugins", "pi-main.pkg", "dist", "index.js") + if bi, err := os.Stat(binAbs); err != nil || bi.Mode()&0o111 == 0 { + t.Errorf("expected main entrypoint executable, err=%v", err) + } +} + +// TestDistributeExtensionNoBin verifies a package without a bin errors. +func TestDistributeExtensionNoBin(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("extension install not supported on windows") + } + t.Setenv("PIGO_HOME", t.TempDir()) + pkg := writePkg(t, `{"name":"pi-nobin","version":"1.0.0"}`, nil) + if _, err := DistributeExtension(pkg, "pi-nobin"); err == nil { + t.Fatal("DistributeExtension without bin = nil error, want error") + } +} diff --git a/pigo/internal/pkgmgr/distribute_theme.go b/pigo/internal/pkgmgr/distribute_theme.go new file mode 100644 index 0000000..0aa5599 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_theme.go @@ -0,0 +1,38 @@ +// This file distributes a classified pi theme into pigo's themes directory +// (#161). Unlike extensions/skills/prompts, pigo has no theme runtime yet, so a +// theme is simply *stored* under $PIGO_HOME/themes// for a future +// consumer — no launcher, no discovery wiring. Storing it (rather than dropping +// it) keeps the install/list/uninstall/update lifecycle uniform: the theme has +// a home, the lockfile records its files, and uninstall can remove it cleanly. +package pkgmgr + +import ( + "fmt" + "os" + "path/filepath" +) + +// DistributeTheme copies the theme package at pkgDir into the themes directory +// under a "/" subdirectory. Because pigo has no theme runtime yet, this +// only stores the theme; there is no discovery step. It returns the absolute +// paths of every file (and the theme dir) it created, for the lockfile. An +// unresolvable themes dir is an error. +func DistributeTheme(pkgDir, name string) ([]string, error) { + themesDir := ThemesDir() + if themesDir == "" { + return nil, fmt.Errorf("pkgmgr: cannot resolve themes dir (PIGO_HOME/home unavailable)") + } + dest := filepath.Join(themesDir, name) + // A stale theme from a prior install must not linger alongside the new one. + if err := os.RemoveAll(dest); err != nil { + return nil, fmt.Errorf("pkgmgr: clear old theme %q: %w", dest, err) + } + if err := os.MkdirAll(themesDir, 0o755); err != nil { + return nil, fmt.Errorf("pkgmgr: create themes dir: %w", err) + } + files, err := copyTree(pkgDir, dest) + if err != nil { + return nil, err + } + return append(files, dest), nil +} diff --git a/pigo/internal/pkgmgr/distribute_theme_test.go b/pigo/internal/pkgmgr/distribute_theme_test.go new file mode 100644 index 0000000..2144da1 --- /dev/null +++ b/pigo/internal/pkgmgr/distribute_theme_test.go @@ -0,0 +1,56 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "testing" +) + +// TestDistributeTheme verifies a theme package is stored under +// $PIGO_HOME/themes// with its files intact. +func TestDistributeTheme(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg := writePkg(t, `{"name":"pi-dark","version":"1.0.0","pi":{"type":"theme"}}`, map[string]string{ + "theme.json": `{"bg":"#000"}`, + "assets/logo.txt": "logo", + }) + + files, err := DistributeTheme(pkg, "pi-dark") + if err != nil { + t.Fatalf("DistributeTheme: %v", err) + } + if _, err := os.Stat(filepath.Join(home, "themes", "pi-dark", "theme.json")); err != nil { + t.Errorf("theme.json not stored: %v", err) + } + if _, err := os.Stat(filepath.Join(home, "themes", "pi-dark", "assets", "logo.txt")); err != nil { + t.Errorf("theme asset not stored: %v", err) + } + if len(files) == 0 { + t.Error("created files empty") + } +} + +// TestDistributeThemeReinstallReplaces verifies reinstall clears stale files. +func TestDistributeThemeReinstallReplaces(t *testing.T) { + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + pkg1 := writePkg(t, `{"name":"pi-t","version":"1.0.0"}`, map[string]string{ + "theme.json": `{}`, + "old.txt": "old", + }) + if _, err := DistributeTheme(pkg1, "pi-t"); err != nil { + t.Fatalf("first install: %v", err) + } + pkg2 := writePkg(t, `{"name":"pi-t","version":"2.0.0"}`, map[string]string{ + "theme.json": `{}`, + }) + if _, err := DistributeTheme(pkg2, "pi-t"); err != nil { + t.Fatalf("reinstall: %v", err) + } + if _, err := os.Stat(filepath.Join(home, "themes", "pi-t", "old.txt")); !os.IsNotExist(err) { + t.Errorf("stale theme file survived reinstall: %v", err) + } +} diff --git a/pigo/internal/pkgmgr/fetch.go b/pigo/internal/pkgmgr/fetch.go new file mode 100644 index 0000000..741711f --- /dev/null +++ b/pigo/internal/pkgmgr/fetch.go @@ -0,0 +1,218 @@ +// This file fetches a pi package's contents from npm (#156). Rather than +// implement an npm registry client, pigo shells out to the user's installed +// `npm` — specifically `npm pack`, which downloads a package as a .tgz tarball +// without running install scripts. pigo then extracts that tarball into a +// temporary directory for the classify/distribute steps that follow. +// +// The fetch is deliberately side-effect-light: `npm pack` neither installs +// dependencies nor runs lifecycle scripts, so downloading a package cannot +// execute its code. Running the extracted extension is a separate, later step. +// +// npm packs every package into a top-level "package/" directory inside the +// tarball; Fetch returns the path to that extracted directory. +package pkgmgr + +import ( + "archive/tar" + "compress/gzip" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" +) + +// FetchResult describes a package fetched to a temporary directory. +type FetchResult struct { + // Dir is the extracted package directory (the tarball's "package/" root). + Dir string + // TempRoot is the temporary directory holding Dir and the tarball; the + // caller must Cleanup it when done. + TempRoot string +} + +// Cleanup removes the temporary directory tree created by Fetch. Safe to call +// on a zero FetchResult (no-op). +func (r FetchResult) Cleanup() error { + if r.TempRoot == "" { + return nil + } + return os.RemoveAll(r.TempRoot) +} + +// npmExecutable is the npm binary name; a variable so tests can stub it. +var npmExecutable = "npm" + +// EnsureNPM reports an actionable error when npm is not on PATH. The install +// command calls this before doing any work so it fails fast with guidance +// rather than deep inside a fetch. +func EnsureNPM() error { + if _, err := exec.LookPath(npmExecutable); err != nil { + return fmt.Errorf("npm not found; install Node.js/npm to use pigo install") + } + return nil +} + +// Fetch downloads the package named by ref using `npm pack` and extracts it into +// a fresh temporary directory. On success the caller owns the returned +// FetchResult and must call Cleanup. On any failure the temporary directory is +// removed before returning, so a failed fetch leaves nothing behind. +// +// npm's own error output (unknown package, network failure, auth) is included +// in the returned error so the user sees why the fetch failed. +func Fetch(ref PackageRef) (FetchResult, error) { + if err := EnsureNPM(); err != nil { + return FetchResult{}, err + } + + tmp, err := os.MkdirTemp("", "pigo-pkg-*") + if err != nil { + return FetchResult{}, fmt.Errorf("pkgmgr: create temp dir: %w", err) + } + // From here on, remove tmp on any error path. + fail := func(e error) (FetchResult, error) { + _ = os.RemoveAll(tmp) + return FetchResult{}, e + } + + spec := ref.Name + if ref.Version != "" { + spec += "@" + ref.Version + } + + // `npm pack ` writes a .tgz into --pack-destination and prints its + // filename. --ignore-scripts guards against packing-time script execution. + cmd := exec.Command(npmExecutable, "pack", spec, + "--pack-destination", tmp, + "--ignore-scripts", + "--loglevel", "error") + var stderr strings.Builder + cmd.Stderr = &stderr + out, err := cmd.Output() + if err != nil { + msg := strings.TrimSpace(stderr.String()) + if msg == "" { + msg = err.Error() + } + return fail(fmt.Errorf("npm pack %s failed: %s", spec, msg)) + } + + tarball, err := locateTarball(tmp, out) + if err != nil { + return fail(err) + } + + dest := filepath.Join(tmp, "extracted") + if err := extractTarGz(tarball, dest); err != nil { + return fail(err) + } + // npm packs into a top-level "package/" directory. + pkgDir := filepath.Join(dest, "package") + if _, err := os.Stat(pkgDir); err != nil { + return fail(fmt.Errorf("pkgmgr: extracted tarball missing package/ dir: %w", err)) + } + return FetchResult{Dir: pkgDir, TempRoot: tmp}, nil +} + +// locateTarball resolves the .tgz path that `npm pack` produced. npm prints the +// tarball filename on stdout; when that is unhelpful we fall back to scanning +// the destination directory for a single .tgz. +func locateTarball(dir string, packStdout []byte) (string, error) { + if name := strings.TrimSpace(string(packStdout)); name != "" { + // npm may print just the filename; join to dir if it isn't absolute. + cand := name + if !filepath.IsAbs(cand) { + cand = filepath.Join(dir, filepath.Base(name)) + } + if _, err := os.Stat(cand); err == nil { + return cand, nil + } + } + entries, err := os.ReadDir(dir) + if err != nil { + return "", fmt.Errorf("pkgmgr: read pack dir: %w", err) + } + for _, e := range entries { + if strings.HasSuffix(e.Name(), ".tgz") { + return filepath.Join(dir, e.Name()), nil + } + } + return "", fmt.Errorf("pkgmgr: npm pack produced no .tgz in %s", dir) +} + +// extractTarGz extracts a gzip-compressed tar archive into dest, creating dest. +// It guards against path traversal (a "../" entry escaping dest) and skips any +// entry that is not a regular file or directory. +func extractTarGz(tarball, dest string) error { + f, err := os.Open(tarball) + if err != nil { + return fmt.Errorf("pkgmgr: open tarball: %w", err) + } + defer f.Close() + + gz, err := gzip.NewReader(f) + if err != nil { + return fmt.Errorf("pkgmgr: gzip reader: %w", err) + } + defer gz.Close() + + if err := os.MkdirAll(dest, 0o755); err != nil { + return fmt.Errorf("pkgmgr: create extract dir: %w", err) + } + + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + if err != nil { + return fmt.Errorf("pkgmgr: read tar: %w", err) + } + target, err := safeJoin(dest, hdr.Name) + if err != nil { + return err + } + switch hdr.Typeflag { + case tar.TypeDir: + if err := os.MkdirAll(target, 0o755); err != nil { + return fmt.Errorf("pkgmgr: mkdir %q: %w", target, err) + } + case tar.TypeReg: + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return fmt.Errorf("pkgmgr: mkdir parent of %q: %w", target, err) + } + if err := writeFile(target, tr, os.FileMode(hdr.Mode)); err != nil { + return err + } + default: + // Skip symlinks, devices, etc. — npm packages are files + dirs. + } + } + return nil +} + +// safeJoin joins name onto base, rejecting any result that escapes base (a +// tarball path-traversal guard). +func safeJoin(base, name string) (string, error) { + target := filepath.Join(base, name) + cleanBase := filepath.Clean(base) + string(os.PathSeparator) + if target != filepath.Clean(base) && !strings.HasPrefix(target, cleanBase) { + return "", fmt.Errorf("pkgmgr: tarball entry %q escapes extract dir", name) + } + return target, nil +} + +// writeFile writes the tar entry body to target with the given mode. +func writeFile(target string, r io.Reader, mode os.FileMode) error { + out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, mode) + if err != nil { + return fmt.Errorf("pkgmgr: create %q: %w", target, err) + } + defer out.Close() + if _, err := io.Copy(out, r); err != nil { + return fmt.Errorf("pkgmgr: write %q: %w", target, err) + } + return nil +} diff --git a/pigo/internal/pkgmgr/fetch_test.go b/pigo/internal/pkgmgr/fetch_test.go new file mode 100644 index 0000000..d76ab80 --- /dev/null +++ b/pigo/internal/pkgmgr/fetch_test.go @@ -0,0 +1,165 @@ +package pkgmgr + +import ( + "archive/tar" + "compress/gzip" + "os" + "path/filepath" + "runtime" + "testing" +) + +// makeTarGz writes a gzip tarball at path containing the given files, each under +// a top-level "package/" dir (mirroring npm pack layout). +func makeTarGz(t *testing.T, path string, files map[string]string) { + t.Helper() + f, err := os.Create(path) + if err != nil { + t.Fatal(err) + } + defer f.Close() + gz := gzip.NewWriter(f) + defer gz.Close() + tw := tar.NewWriter(gz) + defer tw.Close() + for name, body := range files { + hdr := &tar.Header{Name: name, Mode: 0o644, Size: int64(len(body)), Typeflag: tar.TypeReg} + if err := tw.WriteHeader(hdr); err != nil { + t.Fatal(err) + } + if _, err := tw.Write([]byte(body)); err != nil { + t.Fatal(err) + } + } +} + +// TestExtractTarGz verifies a tarball extracts with its files intact. +func TestExtractTarGz(t *testing.T) { + tmp := t.TempDir() + tarball := filepath.Join(tmp, "pkg.tgz") + makeTarGz(t, tarball, map[string]string{ + "package/package.json": `{"name":"x","version":"1.0.0"}`, + "package/index.js": "console.log('hi')", + }) + dest := filepath.Join(tmp, "out") + if err := extractTarGz(tarball, dest); err != nil { + t.Fatalf("extractTarGz: %v", err) + } + got, err := os.ReadFile(filepath.Join(dest, "package", "package.json")) + if err != nil { + t.Fatalf("read extracted: %v", err) + } + if string(got) != `{"name":"x","version":"1.0.0"}` { + t.Errorf("extracted content = %q", got) + } +} + +// TestSafeJoinRejectsTraversal verifies a "../" tarball entry is rejected. +func TestSafeJoinRejectsTraversal(t *testing.T) { + if _, err := safeJoin("/tmp/extract", "../../etc/passwd"); err == nil { + t.Error("safeJoin allowed path traversal, want error") + } + if _, err := safeJoin("/tmp/extract", "package/index.js"); err != nil { + t.Errorf("safeJoin rejected legit path: %v", err) + } +} + +// TestEnsureNPMMissing verifies a clear error when npm is absent. +func TestEnsureNPMMissing(t *testing.T) { + old := npmExecutable + npmExecutable = "definitely-not-a-real-binary-xyz" + defer func() { npmExecutable = old }() + err := EnsureNPM() + if err == nil { + t.Fatal("EnsureNPM with missing npm = nil, want error") + } + if !contains(err.Error(), "npm not found") { + t.Errorf("error = %q, want to mention 'npm not found'", err) + } +} + +// TestFetchWithFakeNPM drives Fetch end-to-end using a fake `npm` on PATH that +// produces a tarball, verifying extraction and cleanup without a real registry. +func TestFetchWithFakeNPM(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("fake npm shell script is POSIX-only") + } + binDir := t.TempDir() + // The fake npm: on `pack`, copy a prebuilt tarball into --pack-destination + // and print its filename, mimicking real npm pack output. + srcTarball := filepath.Join(binDir, "src.tgz") + makeTarGz(t, srcTarball, map[string]string{ + "package/package.json": `{"name":"pi-demo","version":"2.0.0"}`, + }) + fakeNPM := filepath.Join(binDir, "npm") + script := `#!/bin/sh +# args: pack --pack-destination --ignore-scripts --loglevel error +dest="" +prev="" +for a in "$@"; do + if [ "$prev" = "--pack-destination" ]; then dest="$a"; fi + prev="$a" +done +cp "` + srcTarball + `" "$dest/pi-demo-2.0.0.tgz" +echo "pi-demo-2.0.0.tgz" +` + if err := os.WriteFile(fakeNPM, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + + ref, err := ParsePackageRef("npm:pi-demo") + if err != nil { + t.Fatal(err) + } + res, err := Fetch(ref) + if err != nil { + t.Fatalf("Fetch: %v", err) + } + defer res.Cleanup() + + pj, err := os.ReadFile(filepath.Join(res.Dir, "package.json")) + if err != nil { + t.Fatalf("read fetched package.json: %v", err) + } + if !contains(string(pj), `"pi-demo"`) { + t.Errorf("package.json = %q", pj) + } + + // Cleanup removes the temp root. + root := res.TempRoot + if err := res.Cleanup(); err != nil { + t.Fatalf("Cleanup: %v", err) + } + if _, err := os.Stat(root); !os.IsNotExist(err) { + t.Errorf("temp root still exists after Cleanup: %v", err) + } +} + +// TestFetchNPMFailurePropagates verifies npm's error is surfaced and no temp +// dir is left behind. +func TestFetchNPMFailurePropagates(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("fake npm shell script is POSIX-only") + } + binDir := t.TempDir() + fakeNPM := filepath.Join(binDir, "npm") + script := `#!/bin/sh +echo "npm error code E404" >&2 +echo "npm error 404 Not Found - GET https://registry.npmjs.org/nope" >&2 +exit 1 +` + if err := os.WriteFile(fakeNPM, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + + ref, _ := ParsePackageRef("npm:nope") + _, err := Fetch(ref) + if err == nil { + t.Fatal("Fetch of failing npm = nil error, want error") + } + if !contains(err.Error(), "E404") { + t.Errorf("error = %q, want npm stderr included", err) + } +} diff --git a/pigo/internal/pkgmgr/install.go b/pigo/internal/pkgmgr/install.go new file mode 100644 index 0000000..ea77088 --- /dev/null +++ b/pigo/internal/pkgmgr/install.go @@ -0,0 +1,122 @@ +// This file is the install orchestrator (#162): it ties together the pieces the +// earlier issues built — fetch (#156), classify (#157), the per-type +// distributors (#158-#161), and the lockfile (#154) — into the single +// `pigo install npm:` flow. +// +// The flow is: parse the reference, fetch+extract the package to a temp dir, +// classify it into one or more pi types, distribute each type to its target +// directory, then record everything laid down in the lockfile so list/uninstall +// /update can act on it. The temp dir is always cleaned up. +package pkgmgr + +import ( + "fmt" + "io" + "sort" + "strings" +) + +// InstallResult reports what an install did, for the CLI to print. +type InstallResult struct { + Name string + Version string + Types []PackageType + // Files is every path laid down on disk across all distributed types. + Files []string +} + +// Install fetches, classifies, and distributes the package named by rawRef +// (e.g. "npm:pi-mcp-adapter"), then records it in the lockfile at lockfilePath. +// Progress is written to logw when non-nil. It returns a summary of the install. +// +// The install is type-driven: a package classified as several types (e.g. +// extension+skill) is distributed to each corresponding directory. If any +// distribution step fails, the error is returned; already-written files are left +// for the caller/uninstall to reconcile via a re-run (distribution is +// idempotent — each distributor clears its own stale target first). +func Install(rawRef, lockfilePath string, logw io.Writer) (InstallResult, error) { + logf := func(format string, a ...any) { + if logw != nil { + fmt.Fprintf(logw, format, a...) + } + } + + ref, err := ParsePackageRef(rawRef) + if err != nil { + return InstallResult{}, err + } + + logf("Fetching %s ...\n", ref.String()) + fetched, err := Fetch(ref) + if err != nil { + return InstallResult{}, err + } + defer fetched.Cleanup() + + return installFetched(fetched.Dir, ref, lockfilePath, logf) +} + +// installFetched runs the post-fetch half of an install: classify the already +// extracted package at pkgDir, distribute each type, and record the result in +// the lockfile. It is shared by Install (#162) and Update (#164) so both go +// through the same classify→distribute→lockfile path. +func installFetched(pkgDir string, ref PackageRef, lockfilePath string, logf func(string, ...any)) (InstallResult, error) { + name, version, types, err := Classify(pkgDir) + if err != nil { + return InstallResult{}, err + } + logf("Installing %s@%s (%s)\n", name, version, joinTypes(types)) + + var files []string + for _, t := range types { + created, derr := distribute(t, pkgDir, name) + if derr != nil { + return InstallResult{}, derr + } + files = append(files, created...) + } + sort.Strings(files) + + lf, err := Load(lockfilePath) + if err != nil { + return InstallResult{}, err + } + lf.Set(InstalledPackage{ + Name: name, + Source: ref.String(), + Version: version, + Types: types, + Files: files, + }) + if err := lf.Save(); err != nil { + return InstallResult{}, err + } + + return InstallResult{Name: name, Version: version, Types: types, Files: files}, nil +} + +// distribute routes one package type to its distributor. An unknown type is an +// error (Classify should never produce one, but guard anyway). +func distribute(t PackageType, pkgDir, name string) ([]string, error) { + switch t { + case TypeExtension: + return DistributeExtension(pkgDir, name) + case TypeSkill: + return DistributeSkill(pkgDir, name) + case TypePrompt: + return DistributePrompt(pkgDir, name) + case TypeTheme: + return DistributeTheme(pkgDir, name) + default: + return nil, fmt.Errorf("pkgmgr: cannot distribute unknown type %q", t) + } +} + +// joinTypes renders a type slice as a comma-separated string for logging. +func joinTypes(types []PackageType) string { + parts := make([]string, len(types)) + for i, t := range types { + parts[i] = string(t) + } + return strings.Join(parts, ", ") +} diff --git a/pigo/internal/pkgmgr/install_test.go b/pigo/internal/pkgmgr/install_test.go new file mode 100644 index 0000000..3983662 --- /dev/null +++ b/pigo/internal/pkgmgr/install_test.go @@ -0,0 +1,140 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "runtime" + "testing" +) + +// fakeNPMForInstall writes a fake `npm` onto PATH that packs a prebuilt tarball +// (built from files) into --pack-destination, mimicking `npm pack`. +func fakeNPMForInstall(t *testing.T, tarballFiles map[string]string, packName string) { + t.Helper() + binDir := t.TempDir() + srcTarball := filepath.Join(binDir, "src.tgz") + makeTarGz(t, srcTarball, tarballFiles) + script := `#!/bin/sh +dest="" +prev="" +for a in "$@"; do + if [ "$prev" = "--pack-destination" ]; then dest="$a"; fi + prev="$a" +done +cp "` + srcTarball + `" "$dest/` + packName + `" +echo "` + packName + `" +` + if err := os.WriteFile(filepath.Join(binDir, "npm"), []byte(script), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) +} + +// TestInstallExtensionEndToEnd drives Install for an extension package through +// fetch (fake npm) → classify → distribute → lockfile. +func TestInstallExtensionEndToEnd(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("fake npm shell script + extension install are POSIX-only") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + fakeNPMForInstall(t, map[string]string{ + "package/package.json": `{"name":"pi-demo","version":"1.2.0","bin":"./cli.js"}`, + "package/cli.js": "#!/usr/bin/env node\nconsole.log('hi')\n", + }, "pi-demo-1.2.0.tgz") + + lockPath := filepath.Join(home, "packages.json") + res, err := Install("npm:pi-demo", lockPath, nil) + if err != nil { + t.Fatalf("Install: %v", err) + } + if res.Name != "pi-demo" || res.Version != "1.2.0" { + t.Errorf("result = %+v, want pi-demo@1.2.0", res) + } + if len(res.Types) != 1 || res.Types[0] != TypeExtension { + t.Errorf("types = %v, want [extension]", res.Types) + } + + // Launcher exists in plugins. + if _, err := os.Stat(filepath.Join(home, "plugins", "pi-demo")); err != nil { + t.Errorf("launcher not installed: %v", err) + } + + // Lockfile records the package. + lf, err := Load(lockPath) + if err != nil { + t.Fatalf("Load lockfile: %v", err) + } + p, ok := lf.Get("pi-demo") + if !ok { + t.Fatal("lockfile missing pi-demo") + } + if p.Source != "npm:pi-demo" || p.Version != "1.2.0" { + t.Errorf("lockfile entry = %+v", p) + } + if len(p.Files) == 0 { + t.Error("lockfile entry has no files") + } +} + +// TestInstallMultiType drives Install for a package that is both extension and +// skill, verifying both distributions happen and both types are recorded. +func TestInstallMultiType(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX-only") + } + home := t.TempDir() + skills := t.TempDir() + t.Setenv("PIGO_HOME", home) + t.Setenv("PIGO_SKILLS_DIR", skills) + + fakeNPMForInstall(t, map[string]string{ + "package/package.json": `{"name":"combo","version":"1.0.0","bin":"./x.js","pi":{"types":["extension","skill"]}}`, + "package/x.js": "#!/usr/bin/env node\n", + "package/SKILL.md": "---\nname: combo\ndescription: d\n---\nbody", + }, "combo-1.0.0.tgz") + + res, err := Install("npm:combo", filepath.Join(home, "packages.json"), nil) + if err != nil { + t.Fatalf("Install: %v", err) + } + if len(res.Types) != 2 { + t.Errorf("types = %v, want extension+skill", res.Types) + } + if _, err := os.Stat(filepath.Join(home, "plugins", "combo")); err != nil { + t.Errorf("extension launcher missing: %v", err) + } + if _, err := os.Stat(filepath.Join(skills, "combo", "SKILL.md")); err != nil { + t.Errorf("skill not installed: %v", err) + } +} + +// TestInstallUnrecognized verifies a non-pi package fails install with a clear +// error and writes no lockfile entry. +func TestInstallUnrecognized(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX-only") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + + fakeNPMForInstall(t, map[string]string{ + "package/package.json": `{"name":"lodash","version":"4.0.0"}`, + }, "lodash-4.0.0.tgz") + + lockPath := filepath.Join(home, "packages.json") + if _, err := Install("npm:lodash", lockPath, nil); err == nil { + t.Fatal("Install of non-pi package = nil error, want error") + } + if _, err := os.Stat(lockPath); !os.IsNotExist(err) { + t.Errorf("lockfile written for failed install: %v", err) + } +} + +// TestInstallBadRef verifies an invalid reference is rejected before any fetch. +func TestInstallBadRef(t *testing.T) { + if _, err := Install("github:owner/repo", filepath.Join(t.TempDir(), "packages.json"), nil); err == nil { + t.Fatal("Install with non-npm ref = nil error, want error") + } +} diff --git a/pigo/internal/pkgmgr/layout.go b/pigo/internal/pkgmgr/layout.go new file mode 100644 index 0000000..4f42c99 --- /dev/null +++ b/pigo/internal/pkgmgr/layout.go @@ -0,0 +1,110 @@ +// This file defines pigo's install directory layout (#154): where each pi +// package type is placed so pigo's existing discovery mechanisms load it with +// no extra configuration. The paths intentionally match the conventions already +// used elsewhere in cmd/pigo and internal/*: +// +// - extensions → $PIGO_HOME/plugins (internal/plugin.Discover) +// - prompts → $PIGO_HOME/commands (runtime.LoadUserCommandsDir) +// - themes → $PIGO_HOME/themes (no runtime consumer yet) +// - skills → skills dir (~/.agents/skills, PIGO_SKILLS_DIR override) +// +// Skills are the one exception to the $PIGO_HOME root: pigo loads skills from +// ~/.agents/skills (overridable with PIGO_SKILLS_DIR), so SkillsDir honors that +// rather than nesting under $PIGO_HOME. +package pkgmgr + +import ( + "os" + "path/filepath" +) + +// Home returns the pigo home directory: $PIGO_HOME, or ~/.pigo when unset. It +// returns "" when the home directory cannot be resolved and no override is set, +// matching trust.DefaultPath's "unavailable" contract. +func Home() string { + if dir := os.Getenv("PIGO_HOME"); dir != "" { + return dir + } + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".pigo") +} + +// PluginsDir returns $PIGO_HOME/plugins, where installed extensions (including +// MCP adapters) are laid down for internal/plugin.Discover. It returns "" when +// Home is unavailable. +func PluginsDir() string { + h := Home() + if h == "" { + return "" + } + return filepath.Join(h, "plugins") +} + +// CommandsDir returns $PIGO_HOME/commands, the legacy location for installed +// prompt/command templates (still loaded by runtime.LoadUserCommandsDir for +// back-compat). New installs go to PromptsDir. It returns "" when Home is +// unavailable. +func CommandsDir() string { + h := Home() + if h == "" { + return "" + } + return filepath.Join(h, "commands") +} + +// PromptsDir returns $PIGO_HOME/prompts, the pi-aligned location where installed +// prompt templates are laid down for runtime.LoadUserCommandsDir (which loads +// both prompts/ and the legacy commands/). It returns "" when Home is +// unavailable. +func PromptsDir() string { + h := Home() + if h == "" { + return "" + } + return filepath.Join(h, "prompts") +} + +// ThemesDir returns $PIGO_HOME/themes, where installed themes are stored. pigo +// has no theme runtime yet, so this is a holding location for a future consumer. +// It returns "" when Home is unavailable. +func ThemesDir() string { + h := Home() + if h == "" { + return "" + } + return filepath.Join(h, "themes") +} + +// SkillsDir returns the directory installed skills are placed in: PIGO_SKILLS_DIR +// when set, else ~/.agents/skills — matching cmd/pigo's skill loader. It returns +// "" when the home directory cannot be resolved and no override is set. +func SkillsDir() string { + if dir := os.Getenv("PIGO_SKILLS_DIR"); dir != "" { + return dir + } + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".agents", "skills") +} + +// DirForType returns the install directory for a package type, or "" when the +// type is unknown or the underlying home directory is unavailable. +func DirForType(t PackageType) string { + switch t { + case TypeExtension: + return PluginsDir() + case TypePrompt: + return CommandsDir() + case TypeTheme: + return ThemesDir() + case TypeSkill: + return SkillsDir() + default: + return "" + } +} diff --git a/pigo/internal/pkgmgr/layout_test.go b/pigo/internal/pkgmgr/layout_test.go new file mode 100644 index 0000000..58f936f --- /dev/null +++ b/pigo/internal/pkgmgr/layout_test.go @@ -0,0 +1,58 @@ +package pkgmgr + +import ( + "path/filepath" + "testing" +) + +// TestHomeHonorsPIGOHOME verifies Home prefers PIGO_HOME over the default. +func TestHomeHonorsPIGOHOME(t *testing.T) { + t.Setenv("PIGO_HOME", "/custom/pigo") + if got := Home(); got != "/custom/pigo" { + t.Errorf("Home() = %q, want /custom/pigo", got) + } +} + +// TestTypeDirsUnderHome verifies the plugins/commands/themes dirs nest under +// $PIGO_HOME. +func TestTypeDirsUnderHome(t *testing.T) { + t.Setenv("PIGO_HOME", "/custom/pigo") + cases := map[PackageType]string{ + TypeExtension: "/custom/pigo/plugins", + TypePrompt: "/custom/pigo/commands", + TypeTheme: "/custom/pigo/themes", + } + for typ, want := range cases { + if got := DirForType(typ); got != want { + t.Errorf("DirForType(%s) = %q, want %q", typ, got, want) + } + } +} + +// TestSkillsDirHonorsOverride verifies skills use PIGO_SKILLS_DIR, not $PIGO_HOME. +func TestSkillsDirHonorsOverride(t *testing.T) { + t.Setenv("PIGO_SKILLS_DIR", "/custom/skills") + if got := SkillsDir(); got != "/custom/skills" { + t.Errorf("SkillsDir() = %q, want /custom/skills", got) + } + if got := DirForType(TypeSkill); got != "/custom/skills" { + t.Errorf("DirForType(skill) = %q, want /custom/skills", got) + } +} + +// TestSkillsDirDefault verifies skills default to ~/.agents/skills. +func TestSkillsDirDefault(t *testing.T) { + t.Setenv("PIGO_SKILLS_DIR", "") + t.Setenv("HOME", "/home/tester") + want := filepath.Join("/home/tester", ".agents", "skills") + if got := SkillsDir(); got != want { + t.Errorf("SkillsDir() = %q, want %q", got, want) + } +} + +// TestDirForUnknownType verifies an unknown type yields "". +func TestDirForUnknownType(t *testing.T) { + if got := DirForType(PackageType("bogus")); got != "" { + t.Errorf("DirForType(bogus) = %q, want empty", got) + } +} diff --git a/pigo/internal/pkgmgr/lockfile.go b/pigo/internal/pkgmgr/lockfile.go new file mode 100644 index 0000000..56295d2 --- /dev/null +++ b/pigo/internal/pkgmgr/lockfile.go @@ -0,0 +1,183 @@ +// Package pkgmgr implements pigo's pi-package installer state and layout +// (#154). A pi package is an add-on published to npm — an extension (often an +// MCP adapter), a skill, a prompt/command template, or a theme — installed with +// `pigo install npm:`. This package owns two concerns that the install / +// list / uninstall / update commands all build on: +// +// - The lockfile: a JSON record at $PIGO_HOME/packages.json of every installed +// package (name, source, version, types, and the exact files laid down on +// disk). It is the source of truth for list/uninstall/update, so removal and +// upgrade can find and clean up precisely what an install created. +// - The directory layout: where each package type is placed so pigo's existing +// discovery mechanisms pick it up without extra configuration — extensions +// under $PIGO_HOME/plugins, skills under the skills dir, prompts under +// $PIGO_HOME/commands, themes under $PIGO_HOME/themes. +// +// This file defines the lockfile wire types and their load/save, mirroring the +// conventions already used by internal/trust: a missing file is an empty +// lockfile (not an error), while a present-but-malformed file is a hard error so +// a corrupted store is surfaced rather than silently overwritten. +package pkgmgr + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" +) + +// PackageType is one of the pi package kinds pigo can install. A single package +// may declare several (the npm catalog has combined "extensionskill" entries), +// so an InstalledPackage carries a slice of these. +type PackageType string + +const ( + // TypeExtension is an executable extension (including MCP adapters); it is + // laid down under $PIGO_HOME/plugins and discovered by internal/plugin. + TypeExtension PackageType = "extension" + // TypeSkill is a skill bundle placed under the skills directory. + TypeSkill PackageType = "skill" + // TypePrompt is a prompt/command template placed under $PIGO_HOME/commands. + TypePrompt PackageType = "prompt" + // TypeTheme is a theme; pigo has no theme runtime yet, so it is only stored + // under $PIGO_HOME/themes for a future consumer. + TypeTheme PackageType = "theme" +) + +// InstalledPackage is one entry in the lockfile: everything needed to describe, +// upgrade, or remove a package that was installed. +type InstalledPackage struct { + // Name is the package's identifier (the npm package name). + Name string `json:"name"` + // Source is the original install reference, e.g. "npm:pi-mcp-adapter". + Source string `json:"source"` + // Version is the resolved, installed version string. + Version string `json:"version"` + // Types are the pi package kinds this package was classified as (one or more). + Types []PackageType `json:"types"` + // Files are the absolute paths of every file laid down on disk for this + // package, so uninstall/update can remove exactly what was created. + Files []string `json:"files"` +} + +// Lockfile is the on-disk record of all installed packages, keyed by package +// name. The zero value is not usable; obtain one via Load. +type Lockfile struct { + // Version is the lockfile schema version, for forward migration. + Version int `json:"version"` + // Packages maps package name to its installed record. + Packages map[string]InstalledPackage `json:"packages"` + + // path is where Save writes; not serialized. + path string `json:"-"` +} + +// lockfileVersion is the current schema version written by Save. +const lockfileVersion = 1 + +// DefaultLockfilePath returns the lockfile location: $PIGO_HOME/packages.json, +// or ~/.pigo/packages.json when PIGO_HOME is unset. It returns "" when the home +// directory cannot be resolved and no override is set, mirroring +// trust.DefaultPath so the caller can treat the store as unavailable rather than +// guessing a path. +func DefaultLockfilePath() string { + if dir := os.Getenv("PIGO_HOME"); dir != "" { + return filepath.Join(dir, "packages.json") + } + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".pigo", "packages.json") +} + +// Load reads the lockfile at path. A missing file is not an error: it yields an +// empty lockfile whose Save will create the file. A present-but-malformed file +// is a hard error so a corrupted store is surfaced rather than silently +// overwritten. An empty path yields an in-memory-only lockfile (Save is a no-op). +func Load(path string) (*Lockfile, error) { + lf := &Lockfile{ + Version: lockfileVersion, + Packages: make(map[string]InstalledPackage), + path: path, + } + if path == "" { + return lf, nil + } + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return lf, nil // no lockfile yet → empty + } + return nil, fmt.Errorf("pkgmgr: read lockfile %q: %w", path, err) + } + if err := json.Unmarshal(data, lf); err != nil { + return nil, fmt.Errorf("pkgmgr: parse lockfile %q: %w", path, err) + } + if lf.Packages == nil { + lf.Packages = make(map[string]InstalledPackage) + } + lf.path = path + return lf, nil +} + +// Save writes the lockfile to its path as human-readable, indented JSON, with +// package keys in sorted order for a stable diff. Save creates the parent +// directory if needed. It is a no-op when the lockfile has no path (empty-path +// Load), so in-memory use never touches disk. +func (lf *Lockfile) Save() error { + if lf.path == "" { + return nil + } + if lf.Version == 0 { + lf.Version = lockfileVersion + } + if err := os.MkdirAll(filepath.Dir(lf.path), 0o755); err != nil { + return fmt.Errorf("pkgmgr: create lockfile dir: %w", err) + } + data, err := json.MarshalIndent(lf, "", " ") + if err != nil { + return fmt.Errorf("pkgmgr: encode lockfile: %w", err) + } + data = append(data, '\n') + if err := os.WriteFile(lf.path, data, 0o644); err != nil { + return fmt.Errorf("pkgmgr: write lockfile %q: %w", lf.path, err) + } + return nil +} + +// Get returns the installed record for name and whether it exists. +func (lf *Lockfile) Get(name string) (InstalledPackage, bool) { + p, ok := lf.Packages[name] + return p, ok +} + +// Set records (or replaces) a package entry in memory. Call Save to persist. +func (lf *Lockfile) Set(p InstalledPackage) { + if lf.Packages == nil { + lf.Packages = make(map[string]InstalledPackage) + } + lf.Packages[p.Name] = p +} + +// Remove deletes the entry for name in memory, reporting whether it existed. +// Call Save to persist. +func (lf *Lockfile) Remove(name string) bool { + if _, ok := lf.Packages[name]; !ok { + return false + } + delete(lf.Packages, name) + return true +} + +// List returns all installed packages sorted by name, for stable `pigo list` +// output. +func (lf *Lockfile) List() []InstalledPackage { + out := make([]InstalledPackage, 0, len(lf.Packages)) + for _, p := range lf.Packages { + out = append(out, p) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out +} diff --git a/pigo/internal/pkgmgr/lockfile_test.go b/pigo/internal/pkgmgr/lockfile_test.go new file mode 100644 index 0000000..4ce64ec --- /dev/null +++ b/pigo/internal/pkgmgr/lockfile_test.go @@ -0,0 +1,141 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "testing" +) + +// TestLoadMissingFileIsEmpty verifies a missing lockfile yields an empty, +// usable lockfile rather than an error. +func TestLoadMissingFileIsEmpty(t *testing.T) { + path := filepath.Join(t.TempDir(), "packages.json") + lf, err := Load(path) + if err != nil { + t.Fatalf("Load missing file: %v", err) + } + if len(lf.Packages) != 0 { + t.Errorf("expected empty lockfile, got %d packages", len(lf.Packages)) + } + if lf.Version != lockfileVersion { + t.Errorf("version = %d, want %d", lf.Version, lockfileVersion) + } +} + +// TestSaveThenLoadRoundTrips verifies a written lockfile reads back identically. +func TestSaveThenLoadRoundTrips(t *testing.T) { + path := filepath.Join(t.TempDir(), "sub", "packages.json") // sub dir must be created + lf, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + lf.Set(InstalledPackage{ + Name: "pi-mcp-adapter", + Source: "npm:pi-mcp-adapter", + Version: "1.2.3", + Types: []PackageType{TypeExtension}, + Files: []string{"/home/u/.pigo/plugins/pi-mcp-adapter"}, + }) + lf.Set(InstalledPackage{ + Name: "pi-web-access", + Source: "npm:pi-web-access", + Version: "0.1.0", + Types: []PackageType{TypeExtension, TypeSkill}, + Files: []string{"/home/u/.pigo/plugins/pi-web-access"}, + }) + if err := lf.Save(); err != nil { + t.Fatalf("Save: %v", err) + } + + got, err := Load(path) + if err != nil { + t.Fatalf("reload: %v", err) + } + if len(got.Packages) != 2 { + t.Fatalf("reloaded %d packages, want 2", len(got.Packages)) + } + p, ok := got.Get("pi-web-access") + if !ok { + t.Fatal("pi-web-access missing after reload") + } + if p.Version != "0.1.0" || len(p.Types) != 2 { + t.Errorf("pi-web-access = %+v, unexpected", p) + } +} + +// TestSaveIsIndentedJSON verifies the on-disk format is human-readable. +func TestSaveIsIndentedJSON(t *testing.T) { + path := filepath.Join(t.TempDir(), "packages.json") + lf, _ := Load(path) + lf.Set(InstalledPackage{Name: "x", Source: "npm:x", Version: "1", Types: []PackageType{TypeSkill}}) + if err := lf.Save(); err != nil { + t.Fatalf("Save: %v", err) + } + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read: %v", err) + } + if !contains(string(data), "\n ") { + t.Errorf("expected indented JSON, got:\n%s", data) + } +} + +// TestLoadCorruptFileIsError verifies a malformed lockfile is surfaced, never +// silently overwritten. +func TestLoadCorruptFileIsError(t *testing.T) { + path := filepath.Join(t.TempDir(), "packages.json") + if err := os.WriteFile(path, []byte("{not json"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := Load(path); err == nil { + t.Fatal("expected error for corrupt lockfile, got nil") + } +} + +// TestRemove verifies Remove reports existence and deletes the entry. +func TestRemove(t *testing.T) { + lf, _ := Load("") // in-memory + lf.Set(InstalledPackage{Name: "a", Source: "npm:a", Version: "1"}) + if !lf.Remove("a") { + t.Error("Remove(a) = false, want true") + } + if lf.Remove("a") { + t.Error("Remove(a) second time = true, want false") + } + if _, ok := lf.Get("a"); ok { + t.Error("a still present after Remove") + } +} + +// TestListSorted verifies List returns packages ordered by name. +func TestListSorted(t *testing.T) { + lf, _ := Load("") + lf.Set(InstalledPackage{Name: "zebra", Source: "npm:zebra"}) + lf.Set(InstalledPackage{Name: "alpha", Source: "npm:alpha"}) + lf.Set(InstalledPackage{Name: "mango", Source: "npm:mango"}) + got := lf.List() + want := []string{"alpha", "mango", "zebra"} + for i, p := range got { + if p.Name != want[i] { + t.Errorf("List()[%d] = %q, want %q", i, p.Name, want[i]) + } + } +} + +// TestEmptyPathSaveIsNoop verifies an in-memory lockfile never touches disk. +func TestEmptyPathSaveIsNoop(t *testing.T) { + lf, _ := Load("") + lf.Set(InstalledPackage{Name: "a"}) + if err := lf.Save(); err != nil { + t.Errorf("Save on empty-path lockfile: %v", err) + } +} + +func contains(s, sub string) bool { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} diff --git a/pigo/internal/pkgmgr/ref.go b/pigo/internal/pkgmgr/ref.go new file mode 100644 index 0000000..94e393e --- /dev/null +++ b/pigo/internal/pkgmgr/ref.go @@ -0,0 +1,148 @@ +// This file parses and validates the package reference a user passes to +// `pigo install` (#155). pi packages are published to npm, so a reference looks +// like `npm:pi-mcp-adapter`, `npm:@scope/name`, or either form with a version +// suffix (`npm:pi-mcp-adapter@1.2.3`). Parsing is deliberately strict: an +// unsupported source prefix or an invalid npm package name is rejected up front, +// so the install command fails fast with a clear message rather than handing a +// bad name to npm. +// +// Only the `npm:` source is supported this release (PRD Non-Goals exclude +// github:/file:), so Registry is effectively always "npm"; it is kept explicit +// so a future source can be added without changing callers. +package pkgmgr + +import ( + "fmt" + "strings" +) + +// Registry identifies where a package is fetched from. Only npm is supported. +type Registry string + +// RegistryNPM is the npm registry, the only supported source this release. +const RegistryNPM Registry = "npm" + +// PackageRef is a parsed, validated install reference. +type PackageRef struct { + // Registry is the source registry (always RegistryNPM this release). + Registry Registry + // Name is the bare package name, e.g. "pi-mcp-adapter" or "@scope/name". + Name string + // Version is the requested version (the part after '@'), or "" for latest. + Version string + // Raw is the original reference as typed, e.g. "npm:pi-mcp-adapter@1.2.3". + Raw string +} + +// String returns the canonical reference, reconstructed from the parsed parts. +func (r PackageRef) String() string { + s := string(r.Registry) + ":" + r.Name + if r.Version != "" { + s += "@" + r.Version + } + return s +} + +// ParsePackageRef parses an install reference of the form `npm:[@version]`, +// where is a plain (`pi-mcp-adapter`) or scoped (`@scope/name`) npm +// package name. It returns an error when the source prefix is missing or +// unsupported, or when the package name is invalid. +func ParsePackageRef(ref string) (PackageRef, error) { + raw := strings.TrimSpace(ref) + prefix, rest, found := strings.Cut(raw, ":") + if !found || prefix == "" { + return PackageRef{}, fmt.Errorf("unsupported package source, expected npm:") + } + if Registry(prefix) != RegistryNPM { + return PackageRef{}, fmt.Errorf("unsupported package source %q, expected npm:", prefix) + } + if rest == "" { + return PackageRef{}, fmt.Errorf("missing package name, expected npm:") + } + + name, version := splitNameVersion(rest) + if err := validateNPMName(name); err != nil { + return PackageRef{}, err + } + return PackageRef{ + Registry: RegistryNPM, + Name: name, + Version: version, + Raw: raw, + }, nil +} + +// splitNameVersion separates a name from an optional trailing "@version". A +// leading '@' (scoped package) is preserved: the split happens on the LAST '@' +// only when it is not the scope marker, so "@scope/name@1.2.3" yields +// ("@scope/name", "1.2.3") and "@scope/name" yields ("@scope/name", ""). +func splitNameVersion(s string) (name, version string) { + // For a scoped name, the first char is '@' and is not a version separator. + searchFrom := 0 + if strings.HasPrefix(s, "@") { + searchFrom = 1 + } + if idx := strings.LastIndex(s[searchFrom:], "@"); idx >= 0 { + at := searchFrom + idx + return s[:at], s[at+1:] + } + return s, "" +} + +// validateNPMName checks a bare npm package name against the parts of the npm +// naming rules that matter for safely handing the name to the npm CLI: non-empty, +// no whitespace or control characters, no shell-hostile characters, reasonable +// length, and — for scoped names — a well-formed "@scope/name" shape. It does not +// aim to replicate every nuance of npm's validate-npm-package-name; it rejects +// the classes of input that would be unsafe or clearly wrong. +func validateNPMName(name string) error { + if name == "" { + return fmt.Errorf("invalid npm package name: empty") + } + if len(name) > 214 { + return fmt.Errorf("invalid npm package name %q: exceeds 214 characters", name) + } + if name != strings.ToLower(name) { + return fmt.Errorf("invalid npm package name %q: must be lowercase", name) + } + for _, r := range name { + if r <= ' ' || r == '\x7f' { + return fmt.Errorf("invalid npm package name %q: contains whitespace or control character", name) + } + switch r { + case '"', '\'', '\\', '`', '$', '(', ')', '<', '>', '|', ';', '&', '*', '?', '#', '%', '^', '{', '}', '[', ']', ',', '!', '~', '=', '+', ':': + return fmt.Errorf("invalid npm package name %q: contains illegal character %q", name, string(r)) + } + } + + if strings.HasPrefix(name, "@") { + return validateScopedName(name) + } + // Unscoped names may not start with '.' or '_'. + if name[0] == '.' || name[0] == '_' { + return fmt.Errorf("invalid npm package name %q: may not start with '.' or '_'", name) + } + if strings.Contains(name, "/") { + return fmt.Errorf("invalid npm package name %q: only scoped names may contain '/'", name) + } + return nil +} + +// validateScopedName validates the "@scope/name" shape for a scoped package. +func validateScopedName(name string) error { + rest := strings.TrimPrefix(name, "@") + scope, pkg, found := strings.Cut(rest, "/") + if !found { + return fmt.Errorf("invalid scoped package name %q: expected @scope/name", name) + } + if scope == "" || pkg == "" { + return fmt.Errorf("invalid scoped package name %q: empty scope or name", name) + } + if strings.Contains(pkg, "/") { + return fmt.Errorf("invalid scoped package name %q: too many '/' separators", name) + } + if pkg[0] == '.' || pkg[0] == '_' { + return fmt.Errorf("invalid scoped package name %q: name may not start with '.' or '_'", name) + } + return nil +} diff --git a/pigo/internal/pkgmgr/ref_test.go b/pigo/internal/pkgmgr/ref_test.go new file mode 100644 index 0000000..3b0056a --- /dev/null +++ b/pigo/internal/pkgmgr/ref_test.go @@ -0,0 +1,126 @@ +package pkgmgr + +import "testing" + +// TestParsePlainName verifies a plain npm reference parses. +func TestParsePlainName(t *testing.T) { + r, err := ParsePackageRef("npm:pi-mcp-adapter") + if err != nil { + t.Fatalf("ParsePackageRef: %v", err) + } + if r.Registry != RegistryNPM { + t.Errorf("Registry = %q, want npm", r.Registry) + } + if r.Name != "pi-mcp-adapter" { + t.Errorf("Name = %q, want pi-mcp-adapter", r.Name) + } + if r.Version != "" { + t.Errorf("Version = %q, want empty", r.Version) + } +} + +// TestParseScopedName verifies a scoped npm reference parses with the leading +// '@' preserved and no false version split. +func TestParseScopedName(t *testing.T) { + r, err := ParsePackageRef("npm:@gotgenes/pi-subagents") + if err != nil { + t.Fatalf("ParsePackageRef: %v", err) + } + if r.Name != "@gotgenes/pi-subagents" { + t.Errorf("Name = %q, want @gotgenes/pi-subagents", r.Name) + } + if r.Version != "" { + t.Errorf("Version = %q, want empty", r.Version) + } +} + +// TestParseWithVersion verifies the "@version" suffix splits off correctly for +// both plain and scoped names. +func TestParseWithVersion(t *testing.T) { + cases := []struct { + ref, name, version string + }{ + {"npm:pi-mcp-adapter@1.2.3", "pi-mcp-adapter", "1.2.3"}, + {"npm:@scope/name@0.1.0", "@scope/name", "0.1.0"}, + {"npm:@scope/name", "@scope/name", ""}, + {"npm:pkg@latest", "pkg", "latest"}, + } + for _, c := range cases { + r, err := ParsePackageRef(c.ref) + if err != nil { + t.Errorf("ParsePackageRef(%q): %v", c.ref, err) + continue + } + if r.Name != c.name || r.Version != c.version { + t.Errorf("ParsePackageRef(%q) = {%q, %q}, want {%q, %q}", c.ref, r.Name, r.Version, c.name, c.version) + } + } +} + +// TestParseMissingPrefix verifies a reference without npm: is rejected. +func TestParseMissingPrefix(t *testing.T) { + for _, ref := range []string{"pi-mcp-adapter", ""} { + if _, err := ParsePackageRef(ref); err == nil { + t.Errorf("ParsePackageRef(%q) = nil error, want error", ref) + } + } +} + +// TestParseUnsupportedPrefix verifies non-npm sources are rejected (github:/file:). +func TestParseUnsupportedPrefix(t *testing.T) { + for _, ref := range []string{"github:owner/repo", "file:./local", "pypi:foo"} { + if _, err := ParsePackageRef(ref); err == nil { + t.Errorf("ParsePackageRef(%q) = nil error, want error", ref) + } + } +} + +// TestParseInvalidName verifies illegal npm names are rejected. +func TestParseInvalidName(t *testing.T) { + cases := []string{ + "npm:has space", // whitespace + "npm:UPPER", // uppercase + "npm:bad;rm -rf", // shell metacharacters + "npm:.hidden", // leading dot + "npm:_underscore", // leading underscore + "npm:@scope", // scope without name + "npm:@/name", // empty scope + "npm:@scope/", // empty name + "npm:a/b/c", // too many slashes (unscoped with slash) + "npm:", // empty name after prefix + } + for _, ref := range cases { + if _, err := ParsePackageRef(ref); err == nil { + t.Errorf("ParsePackageRef(%q) = nil error, want error", ref) + } + } +} + +// TestRefStringRoundTrips verifies String reconstructs the canonical reference. +func TestRefStringRoundTrips(t *testing.T) { + cases := []string{ + "npm:pi-mcp-adapter", + "npm:pi-mcp-adapter@1.2.3", + "npm:@scope/name@0.1.0", + } + for _, ref := range cases { + r, err := ParsePackageRef(ref) + if err != nil { + t.Fatalf("ParsePackageRef(%q): %v", ref, err) + } + if got := r.String(); got != ref { + t.Errorf("String() = %q, want %q", got, ref) + } + } +} + +// TestParseTrimsWhitespace verifies surrounding whitespace is tolerated. +func TestParseTrimsWhitespace(t *testing.T) { + r, err := ParsePackageRef(" npm:pi-web-access ") + if err != nil { + t.Fatalf("ParsePackageRef: %v", err) + } + if r.Name != "pi-web-access" { + t.Errorf("Name = %q, want pi-web-access", r.Name) + } +} diff --git a/pigo/internal/pkgmgr/uninstall.go b/pigo/internal/pkgmgr/uninstall.go new file mode 100644 index 0000000..7683d3a --- /dev/null +++ b/pigo/internal/pkgmgr/uninstall.go @@ -0,0 +1,72 @@ +// This file implements the list and uninstall operations (#163) on top of the +// lockfile (#154). Both are lightweight lockfile operations that complement the +// install flow (#162): +// +// - Listing just reads the lockfile and returns its entries (sorted by name). +// - Uninstalling removes every file the install laid down (recorded in the +// lockfile entry's Files), then drops the entry and saves. A file that is +// already gone is skipped so a partial prior removal still converges, and +// directory entries are removed with RemoveAll so a package's payload dir +// (e.g. plugins/.pkg) comes out whole. +package pkgmgr + +import ( + "fmt" + "io" + "os" + "sort" +) + +// ListInstalled returns every package recorded in the lockfile at lockfilePath, +// sorted by name. A missing lockfile yields an empty slice (no packages yet), +// mirroring Load's missing-is-empty convention. +func ListInstalled(lockfilePath string) ([]InstalledPackage, error) { + lf, err := Load(lockfilePath) + if err != nil { + return nil, err + } + return lf.List(), nil +} + +// Uninstall removes the package named name: it deletes every file the install +// recorded, then removes the lockfile entry and saves. Progress is written to +// logw when non-nil. Removing a package that is not installed is an error. +// +// Files are removed longest-path-first so a directory entry is deleted after +// its contents; each path is removed with RemoveAll so both plain files and +// payload directories are handled, and an already-absent path is not an error +// (the goal state is "gone"). The lockfile entry is dropped even if some file +// removals were no-ops, so uninstall always converges the record. +func Uninstall(name, lockfilePath string, logw io.Writer) error { + logf := func(format string, a ...any) { + if logw != nil { + fmt.Fprintf(logw, format, a...) + } + } + + lf, err := Load(lockfilePath) + if err != nil { + return err + } + p, ok := lf.Get(name) + if !ok { + return fmt.Errorf("package not installed: %s", name) + } + + // Remove deepest paths first so directory entries are cleared after any + // nested file entries recorded alongside them. + files := append([]string(nil), p.Files...) + sort.Slice(files, func(i, j int) bool { return len(files[i]) > len(files[j]) }) + for _, f := range files { + if err := os.RemoveAll(f); err != nil { + return fmt.Errorf("pkgmgr: remove %q: %w", f, err) + } + } + logf("Removed %d path(s) for %s\n", len(files), name) + + lf.Remove(name) + if err := lf.Save(); err != nil { + return err + } + return nil +} diff --git a/pigo/internal/pkgmgr/uninstall_test.go b/pigo/internal/pkgmgr/uninstall_test.go new file mode 100644 index 0000000..48f918e --- /dev/null +++ b/pigo/internal/pkgmgr/uninstall_test.go @@ -0,0 +1,142 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "testing" +) + +// TestListInstalled verifies ListInstalled returns entries sorted by name and +// an empty slice when no lockfile exists. +func TestListInstalled(t *testing.T) { + home := t.TempDir() + lockPath := filepath.Join(home, "packages.json") + + // No lockfile yet → empty. + got, err := ListInstalled(lockPath) + if err != nil { + t.Fatalf("ListInstalled (empty): %v", err) + } + if len(got) != 0 { + t.Errorf("empty list = %v, want none", got) + } + + // Seed two packages out of order; expect sorted by name. + lf, err := Load(lockPath) + if err != nil { + t.Fatal(err) + } + lf.Set(InstalledPackage{Name: "zeta", Source: "npm:zeta", Version: "1.0.0", Types: []PackageType{TypeSkill}}) + lf.Set(InstalledPackage{Name: "alpha", Source: "npm:alpha", Version: "2.0.0", Types: []PackageType{TypeExtension}}) + if err := lf.Save(); err != nil { + t.Fatal(err) + } + + got, err = ListInstalled(lockPath) + if err != nil { + t.Fatalf("ListInstalled: %v", err) + } + if len(got) != 2 || got[0].Name != "alpha" || got[1].Name != "zeta" { + t.Errorf("list = %+v, want [alpha zeta] sorted", got) + } +} + +// TestUninstallRemovesFilesAndEntry verifies uninstall deletes the recorded +// files (and payload dirs) and drops the lockfile entry. +func TestUninstallRemovesFilesAndEntry(t *testing.T) { + home := t.TempDir() + lockPath := filepath.Join(home, "packages.json") + + // Lay down a payload dir + a file inside it, and a standalone launcher file. + payload := filepath.Join(home, "plugins", "demo.pkg") + if err := os.MkdirAll(payload, 0o755); err != nil { + t.Fatal(err) + } + inner := filepath.Join(payload, "cli.js") + if err := os.WriteFile(inner, []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + launcher := filepath.Join(home, "plugins", "demo") + if err := os.WriteFile(launcher, []byte("#!/bin/sh\n"), 0o755); err != nil { + t.Fatal(err) + } + + lf, err := Load(lockPath) + if err != nil { + t.Fatal(err) + } + lf.Set(InstalledPackage{ + Name: "demo", + Source: "npm:demo", + Version: "1.0.0", + Types: []PackageType{TypeExtension}, + Files: []string{inner, payload, launcher}, + }) + if err := lf.Save(); err != nil { + t.Fatal(err) + } + + if err := Uninstall("demo", lockPath, nil); err != nil { + t.Fatalf("Uninstall: %v", err) + } + + // Files gone. + if _, err := os.Stat(payload); !os.IsNotExist(err) { + t.Errorf("payload dir still present: %v", err) + } + if _, err := os.Stat(launcher); !os.IsNotExist(err) { + t.Errorf("launcher still present: %v", err) + } + // Lockfile entry gone. + lf2, err := Load(lockPath) + if err != nil { + t.Fatal(err) + } + if _, ok := lf2.Get("demo"); ok { + t.Error("lockfile still has demo after uninstall") + } +} + +// TestUninstallMissingFilesSkipped verifies uninstall converges (removes the +// entry) even when some recorded files are already gone. +func TestUninstallMissingFilesSkipped(t *testing.T) { + home := t.TempDir() + lockPath := filepath.Join(home, "packages.json") + + present := filepath.Join(home, "commands", "x.md") + if err := os.MkdirAll(filepath.Dir(present), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(present, []byte("body"), 0o644); err != nil { + t.Fatal(err) + } + missing := filepath.Join(home, "commands", "gone.md") // never created + + lf, err := Load(lockPath) + if err != nil { + t.Fatal(err) + } + lf.Set(InstalledPackage{Name: "cmds", Version: "1.0.0", Types: []PackageType{TypePrompt}, Files: []string{present, missing}}) + if err := lf.Save(); err != nil { + t.Fatal(err) + } + + if err := Uninstall("cmds", lockPath, nil); err != nil { + t.Fatalf("Uninstall with a missing file: %v", err) + } + if _, err := os.Stat(present); !os.IsNotExist(err) { + t.Errorf("present file not removed: %v", err) + } + lf2, _ := Load(lockPath) + if _, ok := lf2.Get("cmds"); ok { + t.Error("entry not removed after uninstall") + } +} + +// TestUninstallNotInstalled verifies uninstalling an unknown package errors. +func TestUninstallNotInstalled(t *testing.T) { + lockPath := filepath.Join(t.TempDir(), "packages.json") + if err := Uninstall("nope", lockPath, nil); err == nil { + t.Fatal("Uninstall of missing package = nil error, want error") + } +} diff --git a/pigo/internal/pkgmgr/update.go b/pigo/internal/pkgmgr/update.go new file mode 100644 index 0000000..30b5bbd --- /dev/null +++ b/pigo/internal/pkgmgr/update.go @@ -0,0 +1,86 @@ +// This file implements the update operation (#164): bring an installed package +// up to the latest version published on npm. Update builds on the pieces the +// earlier issues provided — the lockfile (#154) records what is installed and +// from which source, Fetch (#156) pulls the latest tarball, and the shared +// installFetched path (#162) re-classifies and re-distributes. +// +// The flow per package is: look up its recorded source, fetch the latest +// version, and compare against the installed version. If unchanged, it is a +// no-op ("up to date"). Otherwise the old files are removed and the freshly +// fetched version is distributed and recorded. Fetch/classify happen before any +// removal, so a fetch failure leaves the existing install untouched. +package pkgmgr + +import ( + "fmt" + "io" +) + +// UpdateResult reports what Update did for one package. +type UpdateResult struct { + Name string + // OldVersion is the version that was installed before the update. + OldVersion string + // NewVersion is the version after the update (== OldVersion when no-op). + NewVersion string + // Updated is true when a newer version was fetched and installed. + Updated bool +} + +// Update brings the single installed package name up to the latest version from +// its recorded source. Progress is written to logw when non-nil. Updating a +// package that is not installed is an error. When the installed version is +// already latest, it is a no-op and Updated is false. +// +// The latest tarball is fetched and classified before the old files are +// removed, so a fetch or classify failure leaves the prior install intact. +func Update(name, lockfilePath string, logw io.Writer) (UpdateResult, error) { + logf := func(format string, a ...any) { + if logw != nil { + fmt.Fprintf(logw, format, a...) + } + } + + lf, err := Load(lockfilePath) + if err != nil { + return UpdateResult{}, err + } + p, ok := lf.Get(name) + if !ok { + return UpdateResult{}, fmt.Errorf("package not installed: %s", name) + } + + // Resolve the source to its latest version by dropping any pinned version, + // so `update` always targets the newest published release. + ref, err := ParsePackageRef(p.Source) + if err != nil { + return UpdateResult{}, fmt.Errorf("pkgmgr: bad recorded source for %s: %w", name, err) + } + ref.Version = "" + + logf("Fetching %s ...\n", ref.String()) + fetched, err := Fetch(ref) + if err != nil { + return UpdateResult{}, err // old install untouched + } + defer fetched.Cleanup() + + _, newVersion, _, err := Classify(fetched.Dir) + if err != nil { + return UpdateResult{}, err // old install untouched + } + if newVersion == p.Version { + logf("%s is up to date\n", name) + return UpdateResult{Name: name, OldVersion: p.Version, NewVersion: p.Version, Updated: false}, nil + } + + // Newer version fetched: remove the old files, then distribute the new one. + if err := Uninstall(name, lockfilePath, logw); err != nil { + return UpdateResult{}, err + } + if _, err := installFetched(fetched.Dir, ref, lockfilePath, logf); err != nil { + return UpdateResult{}, err + } + logf("Updated %s %s -> %s\n", name, p.Version, newVersion) + return UpdateResult{Name: name, OldVersion: p.Version, NewVersion: newVersion, Updated: true}, nil +} diff --git a/pigo/internal/pkgmgr/update_test.go b/pigo/internal/pkgmgr/update_test.go new file mode 100644 index 0000000..42cb0ad --- /dev/null +++ b/pigo/internal/pkgmgr/update_test.go @@ -0,0 +1,141 @@ +package pkgmgr + +import ( + "os" + "path/filepath" + "runtime" + "testing" +) + +// fakeNPMVersioned writes a fake `npm` onto PATH that packs whichever tarball is +// currently pointed to by a small "which version" file, so a test can flip the +// version `npm pack` returns between Install and Update. +func fakeNPMVersioned(t *testing.T, tarballs map[string]map[string]string, whichFile string) { + t.Helper() + binDir := t.TempDir() + // Build every tarball once under binDir; the script copies the one named in + // whichFile. + for ver, files := range tarballs { + makeTarGz(t, filepath.Join(binDir, ver+".tgz"), files) + } + // Script reads whichFile -> version, then copies .tgz as its pack + // output and echoes the tarball name (mimicking `npm pack`). + script := `#!/bin/sh +dest="" +prev="" +for a in "$@"; do + if [ "$prev" = "--pack-destination" ]; then dest="$a"; fi + prev="$a" +done +ver=$(cat "` + whichFile + `") +cp "` + binDir + `/$ver.tgz" "$dest/$ver.tgz" +echo "$ver.tgz" +` + if err := os.WriteFile(filepath.Join(binDir, "npm"), []byte(script), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) +} + +// TestUpdateToNewerVersion installs v1.0.0 then updates to v2.0.0, verifying the +// lockfile version is bumped and the new payload is in place. +func TestUpdateToNewerVersion(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("fake npm shell script is POSIX-only") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + whichFile := filepath.Join(t.TempDir(), "which") + + tarballs := map[string]map[string]string{ + "1.0.0": { + "package/package.json": `{"name":"pi-demo","version":"1.0.0","bin":"./cli.js"}`, + "package/cli.js": "#!/usr/bin/env node\n// v1\n", + }, + "2.0.0": { + "package/package.json": `{"name":"pi-demo","version":"2.0.0","bin":"./cli.js"}`, + "package/cli.js": "#!/usr/bin/env node\n// v2\n", + }, + } + fakeNPMVersioned(t, tarballs, whichFile) + + lockPath := filepath.Join(home, "packages.json") + + // Install v1. + if err := os.WriteFile(whichFile, []byte("1.0.0"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := Install("npm:pi-demo", lockPath, nil); err != nil { + t.Fatalf("Install v1: %v", err) + } + + // Flip to v2 and update. + if err := os.WriteFile(whichFile, []byte("2.0.0"), 0o644); err != nil { + t.Fatal(err) + } + res, err := Update("pi-demo", lockPath, nil) + if err != nil { + t.Fatalf("Update: %v", err) + } + if !res.Updated || res.OldVersion != "1.0.0" || res.NewVersion != "2.0.0" { + t.Errorf("update result = %+v, want 1.0.0->2.0.0 updated", res) + } + + lf, err := Load(lockPath) + if err != nil { + t.Fatal(err) + } + p, ok := lf.Get("pi-demo") + if !ok || p.Version != "2.0.0" { + t.Errorf("lockfile version after update = %+v, want 2.0.0", p) + } + // v2 payload content is present. + data, err := os.ReadFile(filepath.Join(home, "plugins", "pi-demo.pkg", "cli.js")) + if err != nil { + t.Fatalf("read updated payload: %v", err) + } + if want := "// v2"; !contains(string(data), want) { + t.Errorf("payload = %q, want to contain %q", string(data), want) + } +} + +// TestUpdateUpToDate installs v1.0.0 and updates against the same version, +// expecting a no-op (Updated false, version unchanged). +func TestUpdateUpToDate(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX-only") + } + home := t.TempDir() + t.Setenv("PIGO_HOME", home) + whichFile := filepath.Join(t.TempDir(), "which") + + tarballs := map[string]map[string]string{ + "1.0.0": { + "package/package.json": `{"name":"pi-demo","version":"1.0.0","bin":"./cli.js"}`, + "package/cli.js": "#!/usr/bin/env node\n", + }, + } + fakeNPMVersioned(t, tarballs, whichFile) + if err := os.WriteFile(whichFile, []byte("1.0.0"), 0o644); err != nil { + t.Fatal(err) + } + + lockPath := filepath.Join(home, "packages.json") + if _, err := Install("npm:pi-demo", lockPath, nil); err != nil { + t.Fatalf("Install: %v", err) + } + res, err := Update("pi-demo", lockPath, nil) + if err != nil { + t.Fatalf("Update: %v", err) + } + if res.Updated { + t.Errorf("update result = %+v, want no-op (Updated false)", res) + } +} + +// TestUpdateNotInstalled verifies updating an unknown package errors. +func TestUpdateNotInstalled(t *testing.T) { + if _, err := Update("nope", filepath.Join(t.TempDir(), "packages.json"), nil); err == nil { + t.Fatal("Update of missing package = nil error, want error") + } +} diff --git a/pigo/internal/plugin/events.go b/pigo/internal/plugin/events.go new file mode 100644 index 0000000..df5ed72 --- /dev/null +++ b/pigo/internal/plugin/events.go @@ -0,0 +1,92 @@ +// This file bridges the agent's event stream to subscribed plugins (US-017, +// #133). The agent loop emits agentcore.AgentEvent values; a plugin declares +// which event types it wants in its manifest. EventNotifier maps each observed +// event to a small, wire-safe payload and hands it to the Manager for +// fire-and-forget delivery — the same "never secrets, only observable fields" +// discipline the stream-json envelope uses. +package plugin + +import ( + "encoding/json" + "io" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// EventNotifier forwards agent lifecycle events to the plugins that subscribed +// to them. It is created once per run and its Handle method is wired as the +// event-stream OnEvent callback. A nil *Manager (no plugins) makes NewNotifier +// return nil, and calling Handle on a nil notifier is a safe no-op — callers can +// wire it unconditionally. +type EventNotifier struct { + mgr *Manager + warnLog io.Writer +} + +// NewEventNotifier returns a notifier over mgr, or nil when mgr is nil or has no +// plugins — so the caller can skip the OnEvent wiring entirely in the common +// no-plugin case. warnLog (when non-nil) receives per-plugin delivery failures. +func NewEventNotifier(mgr *Manager, warnLog io.Writer) *EventNotifier { + if mgr == nil || len(mgr.plugins) == 0 { + return nil + } + return &EventNotifier{mgr: mgr, warnLog: warnLog} +} + +// Handle maps ev to a wire payload and dispatches it to subscribed plugins. It +// is a no-op on a nil notifier and when no plugin subscribes to ev's type, so it +// builds a payload only when someone is listening. Delivery is bounded and +// isolated by the Manager, so Handle never blocks the loop beyond the +// per-plugin event timeout. +func (n *EventNotifier) Handle(ev agentcore.AgentEvent) { + if n == nil { + return + } + t := ev.EventType() + if !n.mgr.Subscribers(t) { + return + } + data, err := json.Marshal(eventPayload(ev)) + if err != nil { + data = nil // deliver the bare type rather than dropping the event + } + n.mgr.DispatchEvent(EventParams{Type: t, Data: data}, n.warnLog) +} + +// eventPayload derives the wire-safe payload for an event: only observable, +// non-secret fields (ids, names, counts, stop reasons, streamed text). It +// mirrors the stream-json envelope's field selection so plugin authors and +// stream consumers see the same shape. +func eventPayload(ev agentcore.AgentEvent) map[string]any { + switch e := ev.(type) { + case agentcore.AgentEndEvent: + return map[string]any{"messageCount": len(e.Messages)} + case agentcore.TurnEndEvent: + p := map[string]any{"stopReason": e.Message.StopReason} + if text := agentcore.ContentToText(e.Message.Content); text != "" { + p["text"] = text + } + if calls := e.Message.ToolCalls(); len(calls) > 0 { + names := make([]string, len(calls)) + for i, c := range calls { + names[i] = c.Name + } + p["toolCalls"] = names + } + return p + case agentcore.ToolExecutionStartEvent: + return map[string]any{"toolCallId": e.ToolCallID, "toolName": e.ToolName} + case agentcore.ToolExecutionEndEvent: + return map[string]any{"toolCallId": e.ToolCallID, "toolName": e.ToolName, "isError": e.IsError} + case agentcore.CompactionEvent: + return map[string]any{ + "reason": e.Reason, + "tokensBefore": e.TokensBefore, + "tokensAfter": e.TokensAfter, + "summarizedCount": e.SummarizedCount, + "keptCount": e.KeptCount, + } + default: + return map[string]any{} + } +} diff --git a/pigo/internal/plugin/events_test.go b/pigo/internal/plugin/events_test.go new file mode 100644 index 0000000..78f6f47 --- /dev/null +++ b/pigo/internal/plugin/events_test.go @@ -0,0 +1,262 @@ +// Tests for plugin lifecycle event subscription and delivery (US-017, #133): +// a plugin declares subscribed event types in its manifest, pigo delivers only +// those via one-way `event` notifications, and a slow/hung plugin is isolated by +// the per-event timeout rather than blocking the caller. +package plugin + +import ( + "encoding/json" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// eventPluginSrc declares a subscription to two event types and appends every +// received `event` notification (as one JSON line: {"type":...,"data":...}) to a +// file whose path is passed via the PIGO_EVENT_LOG env var. It lets the test +// assert exactly which events were delivered, in order. +const eventPluginSrc = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` + Params json.RawMessage ` + "`json:\"params\"`" + ` +} + +func main() { + logPath := os.Getenv("PIGO_EVENT_LOG") + sc := bufio.NewScanner(os.Stdin) + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + w := bufio.NewWriter(os.Stdout) + for sc.Scan() { + var r req + if err := json.Unmarshal(sc.Bytes(), &r); err != nil { + continue + } + switch r.Method { + case "initialize": + reply(w, r.ID, json.RawMessage(` + "`" + `{"name":"watcher","version":"1.0","events":["agent_start","tool_execution_end"]}` + "`" + `)) + case "event": + if logPath != "" { + f, _ := os.OpenFile(logPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644) + if f != nil { + fmt.Fprintf(f, "%s\n", r.Params) + f.Close() + } + } + case "shutdown": + return + } + } +} + +func reply(w *bufio.Writer, id *json.RawMessage, result json.RawMessage) { + if id == nil { + return + } + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() +} +` + +// TestPluginSubscribesReportsManifestEvents checks Subscribes reflects exactly +// the manifest's declared event types. +func TestPluginSubscribesReportsManifestEvents(t *testing.T) { + p := &Plugin{Manifest: Manifest{Name: "w", Events: []string{"agent_start", "tool_execution_end"}}} + if !p.Subscribes("agent_start") { + t.Error("should subscribe to agent_start") + } + if !p.Subscribes("tool_execution_end") { + t.Error("should subscribe to tool_execution_end") + } + if p.Subscribes("turn_end") { + t.Error("should NOT subscribe to unlisted turn_end") + } +} + +// TestEventNotifierDeliversSubscribedOnly runs a real event-logging plugin and +// verifies the notifier delivers a subscribed event but drops an unsubscribed +// one — the full path: manifest events → Subscribers gate → payload → RPC. +func TestEventNotifierDeliversSubscribedOnly(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("shell/exec plugin test is unix-oriented") + } + logPath := filepath.Join(t.TempDir(), "events.log") + t.Setenv("PIGO_EVENT_LOG", logPath) + + bin := buildTestPlugin(t, "watcher", eventPluginSrc) + p, err := Load(bin, nil, os.Stderr) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer p.Close() + + m := &Manager{plugins: []*Plugin{p}} + if !m.Subscribers("agent_start") { + t.Fatal("manager should report a subscriber for agent_start") + } + if m.Subscribers("turn_end") { + t.Fatal("manager should report NO subscriber for turn_end") + } + + n := NewEventNotifier(m, os.Stderr) + if n == nil { + t.Fatal("NewEventNotifier should be non-nil with a subscribing plugin") + } + // Subscribed → delivered. + n.Handle(agentcore.AgentStartEvent{}) + // Unsubscribed → dropped (never written). + n.Handle(agentcore.TurnEndEvent{Message: agentcore.AssistantMessage{}}) + // Subscribed → delivered with a payload. + n.Handle(agentcore.ToolExecutionEndEvent{ToolCallID: "c1", ToolName: "grep", IsError: false}) + + // The plugin writes asynchronously; poll briefly for the two expected lines. + var lines []string + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + data, _ := os.ReadFile(logPath) + lines = splitNonEmpty(string(data)) + if len(lines) >= 2 { + break + } + time.Sleep(20 * time.Millisecond) + } + if len(lines) != 2 { + t.Fatalf("want 2 delivered events, got %d: %q", len(lines), lines) + } + var first EventParams + if err := json.Unmarshal([]byte(lines[0]), &first); err != nil { + t.Fatalf("decode first event: %v", err) + } + if first.Type != "agent_start" { + t.Errorf("first event type = %q, want agent_start", first.Type) + } + var second EventParams + if err := json.Unmarshal([]byte(lines[1]), &second); err != nil { + t.Fatalf("decode second event: %v", err) + } + if second.Type != "tool_execution_end" { + t.Errorf("second event type = %q, want tool_execution_end", second.Type) + } + if !containsField(second.Data, "toolName", "grep") { + t.Errorf("second event data missing toolName=grep: %s", second.Data) + } +} + +// TestNewEventNotifierNilWhenNoSubscribers checks the notifier is nil when there +// are no plugins, so the caller can skip wiring OnEvent entirely. +func TestNewEventNotifierNilWhenNoSubscribers(t *testing.T) { + if NewEventNotifier(nil, os.Stderr) != nil { + t.Error("nil manager should yield nil notifier") + } + if NewEventNotifier(&Manager{}, os.Stderr) != nil { + t.Error("empty manager should yield nil notifier") + } +} + +// TestSendEventTimesOutOnHungPlugin verifies event delivery is bounded: a plugin +// that initializes then stops reading stdin does not block SendEvent beyond the +// timeout — it returns a timeout error instead of hanging. +func TestSendEventTimesOutOnHungPlugin(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("shell/exec plugin test is unix-oriented") + } + // hungPluginSrc initializes, then blocks forever without reading further + // stdin, so the OS pipe buffer fills and a Notify write eventually blocks. + const hungPluginSrc = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "time" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + w := bufio.NewWriter(os.Stdout) + if sc.Scan() { + var r req + json.Unmarshal(sc.Bytes(), &r) + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": r.ID, "result": json.RawMessage(` + "`" + `{"name":"hung","events":["agent_start"]}` + "`" + `)}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() + } + time.Sleep(60 * time.Second) // never reads stdin again +} +` + bin := buildTestPlugin(t, "hung", hungPluginSrc) + p, err := Load(bin, nil, os.Stderr) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer p.Close() + + // A single small event will not fill the pipe; force the write to block by + // sending a large payload many times until SendEvent reports the timeout. The + // key assertion is that SendEvent RETURNS (bounded) rather than hanging. + big := make([]byte, 256*1024) + for i := range big { + big[i] = 'x' + } + payload, _ := json.Marshal(map[string]any{"blob": string(big)}) + + start := time.Now() + var lastErr error + for range 50 { + lastErr = p.SendEvent(EventParams{Type: "agent_start", Data: payload}) + if lastErr != nil { + break + } + if time.Since(start) > 10*time.Second { + t.Fatal("SendEvent never reported a timeout on a hung plugin") + } + } + if lastErr == nil { + t.Fatal("expected a timeout error from a hung plugin, got nil") + } + // The bounded return is the contract; the elapsed time per call is ~eventTimeout. + if elapsed := time.Since(start); elapsed > 12*time.Second { + t.Errorf("SendEvent took too long overall (%s) — not bounded", elapsed) + } +} + +// splitNonEmpty splits s on newlines, dropping empty lines. +func splitNonEmpty(s string) []string { + var out []string + for _, line := range strings.Split(s, "\n") { + if line != "" { + out = append(out, line) + } + } + return out +} + +// containsField reports whether raw JSON object has key == value (string). +func containsField(raw json.RawMessage, key, value string) bool { + var m map[string]any + if err := json.Unmarshal(raw, &m); err != nil { + return false + } + s, ok := m[key].(string) + return ok && s == value +} diff --git a/pigo/internal/plugin/manager.go b/pigo/internal/plugin/manager.go new file mode 100644 index 0000000..f6057f2 --- /dev/null +++ b/pigo/internal/plugin/manager.go @@ -0,0 +1,140 @@ +// This file implements plugin discovery and lifecycle management (US-016, #132): +// finding plugin executables under a config directory, loading each, and +// aggregating their tools. Loading is fault-tolerant — one plugin that fails to +// start or handshake is logged and skipped so the rest still load. +package plugin + +import ( + "fmt" + "io" + "os" + "path/filepath" + "sort" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// Manager owns a set of loaded plugins and their aggregated tools. It is not safe +// for concurrent modification; load once at startup, then read. +type Manager struct { + plugins []*Plugin +} + +// Discover finds and loads every plugin under dir. A plugin is any executable +// regular file directly inside dir (non-executable files and subdirectories are +// ignored). Each plugin is launched and handshaked; a failure is written to +// warnLog (when non-nil) and that plugin is skipped. A missing dir is not an +// error — it yields an empty Manager. pluginStderr, when non-nil, receives every +// plugin's stderr. +func Discover(dir string, warnLog, pluginStderr io.Writer) (*Manager, error) { + m := &Manager{} + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return m, nil // no plugins directory → no plugins + } + return nil, fmt.Errorf("plugin: read dir %q: %w", dir, err) + } + // Deterministic load order for stable tool ordering and diagnostics. + sort.Slice(entries, func(i, j int) bool { return entries[i].Name() < entries[j].Name() }) + + for _, e := range entries { + if e.IsDir() { + continue + } + info, err := e.Info() + if err != nil || !isExecutable(info.Mode()) { + continue + } + path := filepath.Join(dir, e.Name()) + p, err := Load(path, nil, pluginStderr) + if err != nil { + if warnLog != nil { + fmt.Fprintf(warnLog, "pigo: plugin %q failed to load: %v\n", e.Name(), err) + } + continue + } + m.plugins = append(m.plugins, p) + } + return m, nil +} + +// isExecutable reports whether the file mode has any execute bit set. +func isExecutable(mode os.FileMode) bool { + return mode&0o111 != 0 +} + +// Tools returns the aggregated tools of every loaded plugin, in load order. +func (m *Manager) Tools() []agentcore.AgentTool { + var out []agentcore.AgentTool + for _, p := range m.plugins { + out = append(out, p.Tools()...) + } + return out +} + +// Plugins returns the loaded plugins (for command aggregation and diagnostics). +func (m *Manager) Plugins() []*Plugin { return m.plugins } + +// PluginCommand pairs a plugin-declared slash command with the plugin that owns +// it, so a caller can dispatch the command back to its plugin via +// Plugin.CallCommand. +type PluginCommand struct { + // Plugin is the plugin that declared and handles this command. + Plugin *Plugin + // Spec is the command's declaration from the owning plugin's manifest. + Spec CommandSpec +} + +// Commands returns the aggregated slash commands of every loaded plugin, in load +// order (and, within a plugin, in manifest order). Each carries the owning +// plugin so the caller can dispatch it via Plugin.CallCommand. +func (m *Manager) Commands() []PluginCommand { + var out []PluginCommand + for _, p := range m.plugins { + for _, spec := range p.Manifest.Commands { + out = append(out, PluginCommand{Plugin: p, Spec: spec}) + } + } + return out +} + +// Subscribers reports whether any loaded plugin subscribes to the given event +// type. It lets a caller skip building an event payload when nobody is listening +// (US-017, #133). +func (m *Manager) Subscribers(eventType string) bool { + for _, p := range m.plugins { + if p.Subscribes(eventType) { + return true + } + } + return false +} + +// DispatchEvent delivers one lifecycle event to every plugin subscribed to its +// type (US-017, #133). Delivery is best-effort and isolated: each plugin's send +// is bounded by eventTimeout, and a delivery failure to one plugin (timeout, +// dead process) is written to warnLog when non-nil and does not stop delivery to +// the others. It never blocks the agent loop beyond the per-plugin timeout. +func (m *Manager) DispatchEvent(params EventParams, warnLog io.Writer) { + for _, p := range m.plugins { + if !p.Subscribes(params.Type) { + continue + } + if err := p.SendEvent(params); err != nil && warnLog != nil { + fmt.Fprintf(warnLog, "pigo: plugin %q event %q: %v\n", p.Manifest.Name, params.Type, err) + } + } +} + +// Close shuts down every loaded plugin, returning the first error encountered +// (all plugins are attempted regardless). +func (m *Manager) Close() error { + var firstErr error + for _, p := range m.plugins { + if err := p.Close(); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} diff --git a/pigo/internal/plugin/manager_test.go b/pigo/internal/plugin/manager_test.go new file mode 100644 index 0000000..41ba2c4 --- /dev/null +++ b/pigo/internal/plugin/manager_test.go @@ -0,0 +1,137 @@ +// Tests for plugin discovery (US-016, #132): executable detection, deterministic +// order, fault tolerance (a bad plugin is skipped, not fatal), and empty/missing +// directory handling. +package plugin + +import ( + "bytes" + "os" + "path/filepath" + "runtime" + "strings" + "testing" +) + +// TestDiscoverMissingDir checks a missing directory yields an empty manager. +func TestDiscoverMissingDir(t *testing.T) { + m, err := Discover(filepath.Join(t.TempDir(), "nope"), nil, nil) + if err != nil { + t.Fatalf("Discover missing dir: %v", err) + } + if len(m.Plugins()) != 0 { + t.Errorf("want no plugins, got %d", len(m.Plugins())) + } +} + +// TestDiscoverSkipsNonExecutable checks non-executable files are ignored. +func TestDiscoverSkipsNonExecutable(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix executable-bit semantics not applicable on windows") + } + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "readme.txt"), []byte("hi"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.Mkdir(filepath.Join(dir, "subdir"), 0o755); err != nil { + t.Fatal(err) + } + m, err := Discover(dir, nil, nil) + if err != nil { + t.Fatalf("Discover: %v", err) + } + if len(m.Plugins()) != 0 { + t.Errorf("non-executable/dir entries should be skipped, got %d plugins", len(m.Plugins())) + } +} + +// TestDiscoverLoadsAndIsolatesBad checks that a good plugin loads and a bad one +// (executable that isn't a valid plugin) is logged and skipped rather than +// aborting discovery. +func TestDiscoverLoadsAndIsolatesBad(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("shell-script plugins are unix-only in this test") + } + dir := t.TempDir() + + // Good plugin: compiled from echoPluginSrc, placed inside the discovery dir. + good := buildTestPlugin(t, "aaa-echo", echoPluginSrc) + if err := os.Rename(good, filepath.Join(dir, "aaa-echo")); err != nil { + t.Fatal(err) + } + + // Bad plugin: an executable that immediately exits without speaking the + // protocol, so the initialize handshake fails. + bad := filepath.Join(dir, "zzz-bad") + if err := os.WriteFile(bad, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil { + t.Fatal(err) + } + + var warn bytes.Buffer + m, err := Discover(dir, &warn, os.Stderr) + if err != nil { + t.Fatalf("Discover: %v", err) + } + defer m.Close() + + if len(m.Plugins()) != 1 { + t.Fatalf("want 1 good plugin loaded, got %d", len(m.Plugins())) + } + if m.Plugins()[0].Manifest.Name != "echo" { + t.Errorf("loaded wrong plugin: %q", m.Plugins()[0].Manifest.Name) + } + if !strings.Contains(warn.String(), "zzz-bad") { + t.Errorf("bad plugin should be logged, warn=%q", warn.String()) + } + if tools := m.Tools(); len(tools) != 1 || tools[0].Name() != "shout" { + t.Errorf("aggregated tools = %+v, want one 'shout'", tools) + } +} + +// TestManagerCommandsAggregatesInLoadOrder checks that Commands() returns every +// loaded plugin's commands in load order (plugins ordered by discovery, and +// within a plugin by manifest order), each carrying its owning plugin. +func TestManagerCommandsAggregatesInLoadOrder(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("shell-script/compiled plugins are unix-only in this test") + } + dir := t.TempDir() + + // Two command plugins. Load order is by filename, so "a-cmd" loads before + // "b-cmd"; within each plugin the manifest declares greet then bye. + a := buildTestPlugin(t, "a-cmd", cmdPluginSrc) + if err := os.Rename(a, filepath.Join(dir, "a-cmd")); err != nil { + t.Fatal(err) + } + b := buildTestPlugin(t, "b-cmd", cmdPluginSrc) + if err := os.Rename(b, filepath.Join(dir, "b-cmd")); err != nil { + t.Fatal(err) + } + + m, err := Discover(dir, nil, os.Stderr) + if err != nil { + t.Fatalf("Discover: %v", err) + } + defer m.Close() + + cmds := m.Commands() + // Both plugins report the same name "cmd" (from cmdPluginSrc's manifest), so + // the two loaded *Plugin pointers are what distinguish load order. + if len(m.Plugins()) != 2 { + t.Fatalf("want 2 plugins loaded, got %d", len(m.Plugins())) + } + p0, p1 := m.Plugins()[0], m.Plugins()[1] + want := []PluginCommand{ + {Plugin: p0, Spec: CommandSpec{Name: "greet", Description: "greets"}}, + {Plugin: p0, Spec: CommandSpec{Name: "bye", Description: "farewell"}}, + {Plugin: p1, Spec: CommandSpec{Name: "greet", Description: "greets"}}, + {Plugin: p1, Spec: CommandSpec{Name: "bye", Description: "farewell"}}, + } + if len(cmds) != len(want) { + t.Fatalf("Commands() len = %d, want %d (%+v)", len(cmds), len(want), cmds) + } + for i, w := range want { + if cmds[i].Plugin != w.Plugin || cmds[i].Spec != w.Spec { + t.Errorf("Commands()[%d] = {%p, %+v}, want {%p, %+v}", i, cmds[i].Plugin, cmds[i].Spec, w.Plugin, w.Spec) + } + } +} diff --git a/pigo/internal/plugin/manifest.go b/pigo/internal/plugin/manifest.go new file mode 100644 index 0000000..face54a --- /dev/null +++ b/pigo/internal/plugin/manifest.go @@ -0,0 +1,111 @@ +// Package plugin implements pigo's external plugin system (US-016, #132): an +// executable written in any language registers custom tools (and, later, slash +// commands) with pigo without touching pigo's source. pigo launches each plugin +// as a child process and speaks line-delimited JSON-RPC 2.0 over its stdio, +// reusing internal/jsonrpc as the transport. +// +// Protocol (client = pigo, server = plugin): +// +// - initialize → Manifest {name, version, tools[], commands[]} +// The handshake. The plugin declares everything it offers up front. +// - tools/call {name, arguments} → CallResult {content, isError} +// pigo forwards a tool invocation; the plugin runs it and returns the text. +// - event {type, data} (notification) +// pigo pushes a subscribed agent lifecycle event (US-017, #133). One-way, +// fire-and-forget: the plugin never replies and a slow plugin is isolated. +// - shutdown (notification) +// Sent on Close so a well-behaved plugin can exit before stdin EOF. +// +// A plugin that crashes or misbehaves is isolated: its Start failure is logged +// and skipped (other plugins still load), and a tool call against a dead plugin +// returns an error result rather than propagating up. +// +// This file defines the wire types exchanged during the handshake and tool call. +package plugin + +import "encoding/json" + +// Manifest is the plugin's self-description, returned from the initialize call. +type Manifest struct { + // Name identifies the plugin; used to namespace its tools and in diagnostics. + Name string `json:"name"` + // Version is an optional free-form version string for diagnostics. + Version string `json:"version,omitempty"` + // Tools are the tools this plugin registers with the agent. + Tools []ToolSpec `json:"tools,omitempty"` + // Commands are the slash commands this plugin registers. + Commands []CommandSpec `json:"commands,omitempty"` + // Events lists the agent lifecycle event types this plugin subscribes to + // (US-017, #133). pigo delivers only these via one-way `event` notifications; + // an empty list means the plugin observes no events. Valid values are the + // agentcore.Event* discriminants (e.g. "agent_start", "tool_execution_end"). + Events []string `json:"events,omitempty"` +} + +// ToolSpec declares one tool a plugin exposes. Schema is the JSON Schema for the +// tool's arguments, passed through verbatim to the agent's tool registry. +type ToolSpec struct { + Name string `json:"name"` + Description string `json:"description"` + Schema json.RawMessage `json:"schema"` +} + +// CommandSpec declares one slash command a plugin exposes. Prompt is the text +// injected as the next user prompt when the command is invoked (matching the +// declarative-command convention); it may be empty if the plugin handles the +// command by other means. +type CommandSpec struct { + Name string `json:"name"` + Description string `json:"description"` + Prompt string `json:"prompt,omitempty"` +} + +// CallParams is the parameter object for a tools/call request. +type CallParams struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` +} + +// CallResult is the reply to a tools/call request. Content is the tool output as +// text; IsError marks a tool-level failure (distinct from a transport error). +type CallResult struct { + Content string `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// CommandCallParams is the parameter object for a commands/call request. It +// mirrors CallParams' naming (name/arguments) so a plugin can decode tool and +// command invocations with the same conventions. Name is the command's name; +// Args carries its free-form arguments (e.g. the text following the slash +// command) as raw JSON, passed through verbatim to the plugin. +type CommandCallParams struct { + Name string `json:"name"` + Args json.RawMessage `json:"arguments"` +} + +// CommandCallResult is the reply to a commands/call request. Prompt is the text +// injected as the next agent turn (matching the declarative-command +// convention); it may be empty if the command produces no prompt. Notifications +// are messages the plugin asks pigo to surface to the user out of band from the +// prompt. +type CommandCallResult struct { + Prompt string `json:"prompt,omitempty"` + Notifications []CommandNotification `json:"notifications,omitempty"` +} + +// CommandNotification is a single message a command asks pigo to surface to the +// user. Message is the human-readable text; Type is an optional severity or +// category hint (e.g. "info", "warning", "error") that pigo may use to style +// the message. +type CommandNotification struct { + Message string `json:"message"` + Type string `json:"type,omitempty"` +} + +// EventParams is the parameter object for an `event` notification. Type is the +// event discriminant (an agentcore.Event* value); Data carries a small, +// wire-safe payload for that event (never secrets — see plugin.EventData). +type EventParams struct { + Type string `json:"type"` + Data json.RawMessage `json:"data,omitempty"` +} diff --git a/pigo/internal/plugin/manifest_test.go b/pigo/internal/plugin/manifest_test.go new file mode 100644 index 0000000..7a8fa18 --- /dev/null +++ b/pigo/internal/plugin/manifest_test.go @@ -0,0 +1,81 @@ +// Tests for the plugin wire types (#261). These assert the commands/call +// request and result types round-trip through JSON marshal/unmarshal so the +// on-the-wire shape (field names, omitempty behavior) stays stable. +package plugin + +import ( + "encoding/json" + "reflect" + "testing" +) + +func TestCommandCallParamsRoundTrip(t *testing.T) { + want := CommandCallParams{ + Name: "review", + Args: json.RawMessage(`{"path":"foo.go","verbose":true}`), + } + + data, err := json.Marshal(want) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + // Params must name its argument field "arguments" to match CallParams. + var shape struct { + Name string `json:"name"` + Args json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal(data, &shape); err != nil { + t.Fatalf("decode shape: %v", err) + } + if shape.Name != want.Name { + t.Errorf("name = %q, want %q", shape.Name, want.Name) + } + + var got CommandCallParams + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if got.Name != want.Name { + t.Errorf("Name = %q, want %q", got.Name, want.Name) + } + if !json.Valid(got.Args) || string(got.Args) != string(want.Args) { + t.Errorf("Args = %s, want %s", got.Args, want.Args) + } +} + +func TestCommandCallResultRoundTrip(t *testing.T) { + want := CommandCallResult{ + Prompt: "Please summarize the following changes.", + Notifications: []CommandNotification{ + {Message: "loaded 3 files", Type: "info"}, + {Message: "1 file skipped", Type: "warning"}, + {Message: "done"}, + }, + } + + data, err := json.Marshal(want) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + var got CommandCallResult + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if !reflect.DeepEqual(got, want) { + t.Errorf("round-trip mismatch:\n got %+v\n want %+v", got, want) + } +} + +// TestCommandCallResultOmitEmpty confirms empty optional fields drop out of the +// wire form, keeping notifications-free results minimal. +func TestCommandCallResultOmitEmpty(t *testing.T) { + data, err := json.Marshal(CommandCallResult{}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if got := string(data); got != "{}" { + t.Errorf("empty result marshaled to %s, want {}", got) + } +} diff --git a/pigo/internal/plugin/plugin.go b/pigo/internal/plugin/plugin.go new file mode 100644 index 0000000..bbbbdc9 --- /dev/null +++ b/pigo/internal/plugin/plugin.go @@ -0,0 +1,196 @@ +// This file implements the plugin client (US-016, #132): it launches a plugin +// executable, performs the initialize handshake, and exposes the plugin's +// declared tools as agentcore.AgentTool values that forward invocations over +// JSON-RPC. Crash isolation lives here — a call against a plugin whose process +// has died returns an error result, never a panic. +package plugin + +import ( + "context" + "encoding/json" + "fmt" + "io" + "slices" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/jsonrpc" +) + +// initTimeout bounds the initialize handshake so a plugin that never replies +// cannot hang plugin discovery. +const initTimeout = 10 * time.Second + +// eventTimeout bounds one fire-and-forget lifecycle-event delivery (US-017, +// #133). It is short so a slow or hung plugin adds only a small, bounded delay +// per event rather than stalling the agent loop; the event is dropped on +// timeout. +const eventTimeout = 2 * time.Second + +// Plugin is a running plugin: its JSON-RPC client plus the manifest it declared +// during initialize. +type Plugin struct { + Manifest Manifest + client *jsonrpc.Client +} + +// Load starts the plugin executable at path (with optional args) and performs +// the initialize handshake. stderr, when non-nil, receives the plugin's stderr +// for logging. The caller must Close the returned Plugin. +func Load(command string, args []string, stderr io.Writer) (*Plugin, error) { + client, err := jsonrpc.NewClient(jsonrpc.Config{ + Command: command, + Args: args, + Stderr: stderr, + }) + if err != nil { + return nil, fmt.Errorf("plugin: launch %q: %w", command, err) + } + + ctx, cancel := context.WithTimeout(context.Background(), initTimeout) + defer cancel() + + raw, err := client.Call(ctx, "initialize", nil) + if err != nil { + _ = client.Close() + return nil, fmt.Errorf("plugin: initialize %q: %w", command, err) + } + var m Manifest + if err := json.Unmarshal(raw, &m); err != nil { + _ = client.Close() + return nil, fmt.Errorf("plugin: decode manifest from %q: %w", command, err) + } + if m.Name == "" { + _ = client.Close() + return nil, fmt.Errorf("plugin %q: manifest has empty name", command) + } + return &Plugin{Manifest: m, client: client}, nil +} + +// Tools adapts each tool the plugin declared into an agentcore.AgentTool that +// forwards Execute over JSON-RPC. +func (p *Plugin) Tools() []agentcore.AgentTool { + out := make([]agentcore.AgentTool, 0, len(p.Manifest.Tools)) + for _, spec := range p.Manifest.Tools { + out = append(out, &pluginTool{plugin: p, spec: spec}) + } + return out +} + +// Close shuts the plugin down: it sends a best-effort shutdown notification then +// closes the transport (which closes stdin and, if needed, kills the child). The +// shutdown notify is bounded by eventTimeout so a plugin that has stopped reading +// its stdin (whose write pipe is full) cannot make Close block on the transport +// write mutex — Close falls through to client.Close, which kills the child. +func (p *Plugin) Close() error { + done := make(chan struct{}) + go func() { _ = p.client.Notify("shutdown", nil); close(done) }() + select { + case <-done: + case <-time.After(eventTimeout): + } + return p.client.Close() +} + +// call forwards a tool invocation to the plugin and returns its result. +func (p *Plugin) call(ctx context.Context, name string, args json.RawMessage) (CallResult, error) { + raw, err := p.client.Call(ctx, "tools/call", CallParams{Name: name, Arguments: args}) + if err != nil { + return CallResult{}, err + } + var res CallResult + if err := json.Unmarshal(raw, &res); err != nil { + return CallResult{}, fmt.Errorf("plugin %q: decode result for %q: %w", p.Manifest.Name, name, err) + } + return res, nil +} + +// CallCommand forwards a slash-command invocation to the plugin over JSON-RPC +// (commands/call) and returns the plugin's result. args carries the command's +// free-form arguments (passed through verbatim). A transport error (e.g. the +// plugin crashed) or a malformed reply is surfaced as a returned error rather +// than a panic — mirroring how pluginTool.Execute isolates a dead plugin, but +// leaving the caller to decide how to present the failure. +func (p *Plugin) CallCommand(ctx context.Context, name string, args json.RawMessage) (CommandCallResult, error) { + raw, err := p.client.Call(ctx, "commands/call", CommandCallParams{Name: name, Args: args}) + if err != nil { + return CommandCallResult{}, fmt.Errorf("plugin %q: command %q: %w", p.Manifest.Name, name, err) + } + var res CommandCallResult + if err := json.Unmarshal(raw, &res); err != nil { + return CommandCallResult{}, fmt.Errorf("plugin %q: decode command result for %q: %w", p.Manifest.Name, name, err) + } + return res, nil +} + +// Subscribes reports whether the plugin asked to receive the given event type in +// its manifest (US-017, #133). pigo only delivers subscribed events. +func (p *Plugin) Subscribes(eventType string) bool { + return slices.Contains(p.Manifest.Events, eventType) +} + +// SendEvent delivers one lifecycle event to the plugin as a one-way `event` +// notification (US-017, #133). Delivery is fire-and-forget and bounded by +// eventTimeout: the underlying write runs on its own goroutine so a plugin that +// has stopped reading its stdin (a hung or slow plugin) cannot block the agent +// loop — the send is abandoned when the timeout elapses and its error returned. +// The dropped write goroutine ends on its own when the plugin dies or Close +// tears the pipe down. +func (p *Plugin) SendEvent(params EventParams) error { + done := make(chan error, 1) + go func() { done <- p.client.Notify("event", params) }() + select { + case err := <-done: + return err + case <-time.After(eventTimeout): + return fmt.Errorf("plugin %q: event %q delivery timed out after %s", p.Manifest.Name, params.Type, eventTimeout) + } +} + +// pluginTool adapts one plugin-declared tool to the agentcore.AgentTool +// interface. All invocations are forwarded to the owning plugin over RPC. +type pluginTool struct { + plugin *Plugin + spec ToolSpec +} + +// Name implements AgentTool. +func (t *pluginTool) Name() string { return t.spec.Name } + +// Description implements AgentTool. +func (t *pluginTool) Description() string { return t.spec.Description } + +// Schema implements AgentTool. An empty schema declared by the plugin degrades +// to a permissive object schema so registration never fails. +func (t *pluginTool) Schema() json.RawMessage { + if len(t.spec.Schema) == 0 { + return json.RawMessage(`{"type":"object"}`) + } + return t.spec.Schema +} + +// ExecutionMode implements AgentTool. Plugin calls cross a process boundary and +// have unknown side effects, so they run sequentially to be safe. +func (t *pluginTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionSequential +} + +// Execute implements AgentTool by forwarding the call to the plugin process. A +// transport error (e.g. the plugin crashed) is isolated: it degrades to an error +// result so a dead plugin cannot take down the agent loop. +func (t *pluginTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + res, err := t.plugin.call(ctx, t.spec.Name, args) + if err != nil { + return agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent( + fmt.Sprintf("%s: plugin call failed: %v", t.spec.Name, err))}, + }, nil + } + result := agentcore.AgentToolResult{ + Content: agentcore.ContentList{agentcore.NewTextContent(res.Content)}, + } + if res.IsError { + result.Details = map[string]any{"isError": true} + } + return result, nil +} diff --git a/pigo/internal/plugin/plugin_test.go b/pigo/internal/plugin/plugin_test.go new file mode 100644 index 0000000..a62f70c --- /dev/null +++ b/pigo/internal/plugin/plugin_test.go @@ -0,0 +1,320 @@ +// Tests for the plugin system (US-016, #132). A test plugin is a tiny Go program +// compiled once per test run; it speaks the JSON-RPC protocol over stdio so the +// tests exercise the real subprocess transport, handshake, tool forwarding, and +// crash isolation — no network or mocks. +package plugin + +import ( + "context" + "encoding/json" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// buildTestPlugin compiles the given Go source into an executable under a temp +// dir and returns its path. The source is a standalone main package. +func buildTestPlugin(t *testing.T, name, src string) string { + t.Helper() + dir := t.TempDir() + srcPath := filepath.Join(dir, name+".go") + if err := os.WriteFile(srcPath, []byte(src), 0o644); err != nil { + t.Fatalf("write plugin source: %v", err) + } + bin := filepath.Join(dir, name) + if runtime.GOOS == "windows" { + bin += ".exe" + } + cmd := exec.Command("go", "build", "-o", bin, srcPath) + cmd.Env = os.Environ() + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("build test plugin: %v\n%s", err, out) + } + return bin +} + +// echoPluginSrc is a plugin that declares one "shout" tool which uppercases its +// "text" argument. +const echoPluginSrc = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "strings" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` + Params json.RawMessage ` + "`json:\"params\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + w := bufio.NewWriter(os.Stdout) + for sc.Scan() { + var r req + if err := json.Unmarshal(sc.Bytes(), &r); err != nil { + continue + } + switch r.Method { + case "initialize": + reply(w, r.ID, json.RawMessage(` + "`" + `{"name":"echo","version":"1.0","tools":[{"name":"shout","description":"uppercase text","schema":{"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}}]}` + "`" + `)) + case "tools/call": + var p struct { + Name string ` + "`json:\"name\"`" + ` + Arguments json.RawMessage ` + "`json:\"arguments\"`" + ` + } + json.Unmarshal(r.Params, &p) + var a struct{ Text string ` + "`json:\"text\"`" + ` } + json.Unmarshal(p.Arguments, &a) + res, _ := json.Marshal(map[string]any{"content": strings.ToUpper(a.Text)}) + reply(w, r.ID, res) + case "shutdown": + return + } + } +} + +func reply(w *bufio.Writer, id *json.RawMessage, result json.RawMessage) { + if id == nil { + return + } + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() +} +` + +// TestPluginLoadAndCall exercises the full path: build → load (handshake) → +// adapt tool → call → result. +func TestPluginLoadAndCall(t *testing.T) { + bin := buildTestPlugin(t, "echo", echoPluginSrc) + p, err := Load(bin, nil, os.Stderr) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer p.Close() + + if p.Manifest.Name != "echo" { + t.Errorf("manifest name = %q, want echo", p.Manifest.Name) + } + tools := p.Tools() + if len(tools) != 1 || tools[0].Name() != "shout" { + t.Fatalf("tools = %+v, want one 'shout'", tools) + } + if tools[0].ExecutionMode() != agentcore.ToolExecutionSequential { + t.Errorf("plugin tool should be sequential") + } + + res, err := tools[0].Execute(context.Background(), "c1", json.RawMessage(`{"text":"hello"}`), nil) + if err != nil { + t.Fatalf("Execute Go error: %v", err) + } + if txt := agentcore.ContentToText(res.Content); txt != "HELLO" { + t.Errorf("result = %q, want HELLO", txt) + } +} + +// crashPluginSrc initializes fine but exits abruptly on the first tools/call. +const crashPluginSrc = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + w := bufio.NewWriter(os.Stdout) + for sc.Scan() { + var r req + if err := json.Unmarshal(sc.Bytes(), &r); err != nil { + continue + } + switch r.Method { + case "initialize": + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": r.ID, "result": json.RawMessage(` + "`" + `{"name":"crash","tools":[{"name":"boom","description":"crashes"}]}` + "`" + `)}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() + case "tools/call": + os.Exit(1) // crash mid-call: no response is ever sent + } + } +} +` + +// TestPluginCrashIsolation checks that a plugin crashing during a tool call +// degrades to an error result, not a panic or a hang. +func TestPluginCrashIsolation(t *testing.T) { + bin := buildTestPlugin(t, "crash", crashPluginSrc) + p, err := Load(bin, nil, os.Stderr) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer p.Close() + + tools := p.Tools() + if len(tools) != 1 { + t.Fatalf("want one tool, got %d", len(tools)) + } + res, err := tools[0].Execute(context.Background(), "c1", json.RawMessage(`{}`), nil) + if err != nil { + t.Fatalf("Execute must not return a Go error even on crash: %v", err) + } + txt := agentcore.ContentToText(res.Content) + if !strings.Contains(txt, "plugin call failed") { + t.Errorf("expected isolated error result, got %q", txt) + } +} + +// cmdPluginSrc declares two slash commands and answers commands/call by echoing +// the command name back as a prompt plus one notification. Its manifest command +// order (greet, then bye) lets tests assert manifest-order aggregation. +const cmdPluginSrc = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` + Params json.RawMessage ` + "`json:\"params\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + w := bufio.NewWriter(os.Stdout) + for sc.Scan() { + var r req + if err := json.Unmarshal(sc.Bytes(), &r); err != nil { + continue + } + switch r.Method { + case "initialize": + reply(w, r.ID, json.RawMessage(` + "`" + `{"name":"cmd","commands":[{"name":"greet","description":"greets"},{"name":"bye","description":"farewell"}]}` + "`" + `)) + case "commands/call": + var p struct { + Name string ` + "`json:\"name\"`" + ` + Args json.RawMessage ` + "`json:\"arguments\"`" + ` + } + json.Unmarshal(r.Params, &p) + res, _ := json.Marshal(map[string]any{ + "prompt": "did:" + p.Name, + "notifications": []map[string]any{ + {"message": "ran " + p.Name, "type": "info"}, + }, + }) + reply(w, r.ID, res) + case "shutdown": + return + } + } +} + +func reply(w *bufio.Writer, id *json.RawMessage, result json.RawMessage) { + if id == nil { + return + } + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, "result": result}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() +} +` + +// TestPluginCallCommand checks CallCommand round-trips a prompt and its +// notifications from a plugin over the real subprocess transport. +func TestPluginCallCommand(t *testing.T) { + bin := buildTestPlugin(t, "cmd", cmdPluginSrc) + p, err := Load(bin, nil, os.Stderr) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer p.Close() + + res, err := p.CallCommand(context.Background(), "greet", json.RawMessage(`{"text":"hi"}`)) + if err != nil { + t.Fatalf("CallCommand: %v", err) + } + if res.Prompt != "did:greet" { + t.Errorf("prompt = %q, want did:greet", res.Prompt) + } + if len(res.Notifications) != 1 { + t.Fatalf("notifications = %+v, want one", res.Notifications) + } + if res.Notifications[0].Message != "ran greet" || res.Notifications[0].Type != "info" { + t.Errorf("notification = %+v, want {ran greet, info}", res.Notifications[0]) + } +} + +// TestPluginCallCommandTransportError checks that a transport error (the plugin +// crashed mid-call) surfaces as a returned error, never a panic. +func TestPluginCallCommandTransportError(t *testing.T) { + bin := buildTestPlugin(t, "cmdcrash", cmdCrashPluginSrc) + p, err := Load(bin, nil, os.Stderr) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer p.Close() + + _, err = p.CallCommand(context.Background(), "greet", json.RawMessage(`{}`)) + if err == nil { + t.Fatal("CallCommand must return an error when the plugin crashes mid-call") + } +} + +// cmdCrashPluginSrc initializes with one command then exits abruptly on the +// first commands/call, so no response is ever sent. +const cmdCrashPluginSrc = `package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" +) + +type req struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + w := bufio.NewWriter(os.Stdout) + for sc.Scan() { + var r req + if err := json.Unmarshal(sc.Bytes(), &r); err != nil { + continue + } + switch r.Method { + case "initialize": + out, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": r.ID, "result": json.RawMessage(` + "`" + `{"name":"cmdcrash","commands":[{"name":"greet","description":"greets"}]}` + "`" + `)}) + fmt.Fprintf(w, "%s\n", out) + w.Flush() + case "commands/call": + os.Exit(1) // crash mid-call: no response is ever sent + } + } +} +` diff --git a/pigo/internal/provider/anthropic.go b/pigo/internal/provider/anthropic.go new file mode 100644 index 0000000..26f0509 --- /dev/null +++ b/pigo/internal/provider/anthropic.go @@ -0,0 +1,334 @@ +// This file implements the Anthropic Messages API streaming decoder (US-008). +// It is a stateful Decoder (see transport.go) that translates Anthropic SSE +// event payloads into the provider-agnostic AssistantMessageEvent set, +// accumulating a partial AssistantMessage as deltas arrive. +// +// Anthropic streams a fixed event sequence: +// +// message_start → seeds id/model and initial usage (input tokens) +// content_block_start → opens a text / thinking / tool_use block at an index +// content_block_delta → text_delta / thinking_delta / signature_delta / +// input_json_delta append to the open block +// content_block_stop → closes the block (tool_use JSON is parsed here) +// message_delta → carries the final stop_reason and output-token usage +// message_stop → terminal; the accumulated message is emitted as done +// error → a runtime error payload → terminal error event +// +// Per the dual failure model (FR-13) the decoder never panics: malformed +// payloads and Anthropic `error` events are surfaced as a returned error (which +// the transport turns into a terminal StreamErrorEvent) rather than crashing. +package provider + +import ( + "encoding/json" + "fmt" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// anthropicBlock accumulates one content block's streaming state, keyed by its +// Anthropic content-block index. text/thinking append their deltas; tool_use +// accumulates a partial JSON string parsed lazily when the block is realized. +type anthropicBlock struct { + kind string // "text" | "thinking" | "tool_use" | "redacted_thinking" + text strings.Builder + thinking strings.Builder + thinkingSig string + textSig string + toolID string + toolName string + toolJSON strings.Builder + redacted bool +} + +// AnthropicDecoder is the stateful SSE decoder for the Anthropic Messages API. +// It implements the transport Decoder interface. It is not safe for concurrent +// use — the transport drives it from a single goroutine. +type AnthropicDecoder struct { + blocks map[int]*anthropicBlock + order []int // content-block indices in first-seen order + + responseID string + responseModel string + inputTokens int + outputTokens int + stopReason string // mapped pigo stop reason (empty until message_delta) + done bool // message_stop / done already emitted +} + +// NewAnthropicDecoder builds a fresh decoder for one streamed response. +func NewAnthropicDecoder() *AnthropicDecoder { + return &AnthropicDecoder{blocks: make(map[int]*anthropicBlock)} +} + +// anthropicEvent is the discriminated envelope shared by every Anthropic SSE +// data payload; fields are populated selectively by event type. +type anthropicEvent struct { + Type string `json:"type"` + Index int `json:"index"` + + Message *struct { + ID string `json:"id"` + Model string `json:"model"` + Usage *anthropicUsage `json:"usage"` + } `json:"message"` + + ContentBlock *struct { + Type string `json:"type"` + // text + Text string `json:"text"` + // thinking + Thinking string `json:"thinking"` + // tool_use + ID string `json:"id"` + Name string `json:"name"` + } `json:"content_block"` + + Delta *struct { + Type string `json:"type"` + // text_delta + Text string `json:"text"` + // thinking_delta + Thinking string `json:"thinking"` + // signature_delta + Signature string `json:"signature"` + // input_json_delta + PartialJSON string `json:"partial_json"` + // message_delta + StopReason string `json:"stop_reason"` + } `json:"delta"` + + Usage *anthropicUsage `json:"usage"` + + Error *struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` +} + +type anthropicUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` +} + +// Decode turns one Anthropic SSE data payload into zero or more StreamEvents. +func (d *AnthropicDecoder) Decode(payload []byte) ([]StreamEvent, error) { + var ev anthropicEvent + if err := json.Unmarshal(payload, &ev); err != nil { + return nil, fmt.Errorf("anthropic: parse event: %w", err) + } + + switch ev.Type { + case "message_start": + return d.onMessageStart(ev), nil + case "content_block_start": + return d.onBlockStart(ev), nil + case "content_block_delta": + return d.onBlockDelta(ev), nil + case "content_block_stop": + // Nothing to emit on stop; the block is already reflected in the partial. + return nil, nil + case "message_delta": + return d.onMessageDelta(ev), nil + case "message_stop": + return d.finishDone(), nil + case "ping": + return nil, nil + case "error": + msg := "anthropic stream error" + if ev.Error != nil { + if ev.Error.Type != "" { + msg = "anthropic " + ev.Error.Type + } + if ev.Error.Message != "" { + msg += ": " + ev.Error.Message + } + } + return nil, fmt.Errorf("%s", msg) + default: + // Unknown event types are ignored (forward-compatible). + return nil, nil + } +} + +// Finish flushes a terminal done event if the stream ended without an explicit +// message_stop (e.g. a clean EOF mid-stream), so a partial response is still +// delivered rather than lost. +func (d *AnthropicDecoder) Finish() ([]StreamEvent, error) { + if d.done { + return nil, nil + } + return d.finishDone(), nil +} + +func (d *AnthropicDecoder) onMessageStart(ev anthropicEvent) []StreamEvent { + if ev.Message != nil { + d.responseID = ev.Message.ID + d.responseModel = ev.Message.Model + if ev.Message.Usage != nil { + d.inputTokens = ev.Message.Usage.InputTokens + d.outputTokens = ev.Message.Usage.OutputTokens + } + } + return []StreamEvent{StreamStartEvent{Partial: d.partial()}} +} + +func (d *AnthropicDecoder) onBlockStart(ev anthropicEvent) []StreamEvent { + if ev.ContentBlock == nil { + return nil + } + b := &anthropicBlock{kind: ev.ContentBlock.Type} + switch ev.ContentBlock.Type { + case "text": + b.text.WriteString(ev.ContentBlock.Text) + case "thinking": + b.thinking.WriteString(ev.ContentBlock.Thinking) + case "redacted_thinking": + b.redacted = true + case "tool_use": + b.toolID = ev.ContentBlock.ID + b.toolName = ev.ContentBlock.Name + } + d.putBlock(ev.Index, b) + + switch ev.ContentBlock.Type { + case "thinking", "redacted_thinking": + return []StreamEvent{StreamThinkingEvent{Partial: d.partial()}} + case "tool_use": + return []StreamEvent{StreamToolCallEvent{Partial: d.partial()}} + default: + return []StreamEvent{StreamTextEvent{Partial: d.partial()}} + } +} + +func (d *AnthropicDecoder) onBlockDelta(ev anthropicEvent) []StreamEvent { + if ev.Delta == nil { + return nil + } + b := d.blocks[ev.Index] + if b == nil { + // A delta for an unseen index: open a bare block so we don't drop data. + b = &anthropicBlock{} + d.putBlock(ev.Index, b) + } + switch ev.Delta.Type { + case "text_delta": + b.text.WriteString(ev.Delta.Text) + return []StreamEvent{StreamTextEvent{Partial: d.partial()}} + case "thinking_delta": + b.thinking.WriteString(ev.Delta.Thinking) + return []StreamEvent{StreamThinkingEvent{Partial: d.partial()}} + case "signature_delta": + // Signature belongs to the thinking block it rides on. + b.thinkingSig += ev.Delta.Signature + return []StreamEvent{StreamThinkingEvent{Partial: d.partial()}} + case "input_json_delta": + b.toolJSON.WriteString(ev.Delta.PartialJSON) + return []StreamEvent{StreamToolCallEvent{Partial: d.partial()}} + default: + return nil + } +} + +func (d *AnthropicDecoder) onMessageDelta(ev anthropicEvent) []StreamEvent { + if ev.Delta != nil && ev.Delta.StopReason != "" { + d.stopReason = mapAnthropicStopReason(ev.Delta.StopReason) + } + if ev.Usage != nil { + // message_delta reports cumulative output tokens (and sometimes input). + if ev.Usage.OutputTokens != 0 { + d.outputTokens = ev.Usage.OutputTokens + } + if ev.Usage.InputTokens != 0 { + d.inputTokens = ev.Usage.InputTokens + } + } + // No standalone event kind for usage/stop-reason accumulation; the values + // surface in the terminal done message. + return nil +} + +// finishDone builds the terminal assistant message and marks the decoder done. +func (d *AnthropicDecoder) finishDone() []StreamEvent { + if d.done { + return nil + } + d.done = true + msg := d.partial() + if msg.StopReason == "" { + msg.StopReason = agentcore.StopReasonEndTurn + } + return []StreamEvent{StreamDoneEvent{Message: msg}} +} + +// putBlock records a block at index, tracking first-seen order. +func (d *AnthropicDecoder) putBlock(index int, b *anthropicBlock) { + if _, seen := d.blocks[index]; !seen { + d.order = append(d.order, index) + } + d.blocks[index] = b +} + +// partial materializes the accumulated state into an AssistantMessage. Content +// blocks are emitted in content-block index order. Tool-use JSON that has not +// yet parsed cleanly is passed through as-is (raw partial), which is valid for +// a still-streaming partial and finalized once the block completes. +func (d *AnthropicDecoder) partial() agentcore.AssistantMessage { + msg := agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + API: "anthropic", + Provider: "anthropic", + StopReason: d.stopReason, + ResponseID: d.responseID, + ResponseModel: d.responseModel, + } + if d.inputTokens != 0 || d.outputTokens != 0 { + msg.Usage = &agentcore.Usage{InputTokens: d.inputTokens, OutputTokens: d.outputTokens} + } + + idx := make([]int, len(d.order)) + copy(idx, d.order) + sort.Ints(idx) + + for _, i := range idx { + b := d.blocks[i] + if b == nil { + continue + } + switch b.kind { + case "thinking", "redacted_thinking": + tc := agentcore.NewThinkingContent(b.thinking.String()) + tc.ThinkingSignature = b.thinkingSig + tc.Redacted = b.redacted + msg.Content = append(msg.Content, tc) + case "tool_use": + args := json.RawMessage(strings.TrimSpace(b.toolJSON.String())) + if len(args) == 0 { + args = json.RawMessage("{}") + } + msg.Content = append(msg.Content, agentcore.NewToolCallContent(b.toolID, b.toolName, args)) + default: // text + tc := agentcore.NewTextContent(b.text.String()) + tc.TextSignature = b.textSig + msg.Content = append(msg.Content, tc) + } + } + return msg +} + +// mapAnthropicStopReason maps an Anthropic stop_reason to the pigo StopReason +// set. Unknown reasons default to end_turn (a natural, non-error stop). +func mapAnthropicStopReason(reason string) string { + switch reason { + case "max_tokens": + return agentcore.StopReasonLength + case "tool_use": + return agentcore.StopReasonToolUse + case "end_turn", "stop_sequence": + return agentcore.StopReasonEndTurn + default: + return agentcore.StopReasonEndTurn + } +} diff --git a/pigo/internal/provider/anthropic_test.go b/pigo/internal/provider/anthropic_test.go new file mode 100644 index 0000000..b727e1a --- /dev/null +++ b/pigo/internal/provider/anthropic_test.go @@ -0,0 +1,299 @@ +package provider + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// A recorded Anthropic Messages API SSE stream covering a text block, a +// thinking block (with signature), and a tool_use block, ending with a +// tool_use stop reason and output-token usage. Trimmed but structurally +// faithful to the real wire format. +const anthropicToolUseSSE = `event: message_start +data: {"type":"message_start","message":{"id":"msg_01ABC","model":"claude-opus-4-8","usage":{"input_tokens":42,"output_tokens":1}}} + +event: content_block_start +data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"Let me think. "}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"Use the tool."}} + +event: content_block_delta +data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sigABC"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":0} + +event: content_block_start +data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"I'll check "}} + +event: content_block_delta +data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"the weather."}} + +event: content_block_stop +data: {"type":"content_block_stop","index":1} + +event: content_block_start +data: {"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"toolu_01","name":"get_weather"}} + +event: content_block_delta +data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"city\":"}} + +event: content_block_delta +data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":" \"SF\"}"}} + +event: content_block_stop +data: {"type":"content_block_stop","index":2} + +event: message_delta +data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":57}} + +event: message_stop +data: {"type":"message_stop"} + +` + +// feedSSE splits a recorded SSE body into events (blank-line separated), +// extracts each `data:` payload, and drives the decoder exactly as the +// transport pump would, returning all emitted events plus the final message. +func feedSSE(t *testing.T, dec Decoder, body string) ([]StreamEvent, agentcore.AssistantMessage) { + t.Helper() + var events []StreamEvent + for _, block := range strings.Split(body, "\n\n") { + var payload strings.Builder + for _, line := range strings.Split(block, "\n") { + line = strings.TrimRight(line, "\r") + if strings.HasPrefix(line, "data:") { + payload.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + } + } + if payload.Len() == 0 { + continue + } + // [DONE] is a transport-level terminator, not a decoder payload. + if payload.String() == "[DONE]" { + continue + } + evs, err := dec.Decode([]byte(payload.String())) + if err != nil { + t.Fatalf("decode %q: %v", payload.String(), err) + } + events = append(events, evs...) + } + finalEvents, err := dec.Finish() + if err != nil { + t.Fatalf("finish: %v", err) + } + events = append(events, finalEvents...) + + var final agentcore.AssistantMessage + for _, ev := range events { + if d, ok := ev.(StreamDoneEvent); ok { + final = d.Message + } + } + return events, final +} + +func TestAnthropicDecoderToolUseStream(t *testing.T) { + dec := NewAnthropicDecoder() + events, final := feedSSE(t, dec, anthropicToolUseSSE) + + // The first emitted event must be a start event. + if len(events) == 0 || events[0].EventKind() != StreamEventStart { + t.Fatalf("expected a start event first, got %v", eventKinds(events)) + } + // The last emitted event must be the terminal done event. + if events[len(events)-1].EventKind() != StreamEventDone { + t.Fatalf("expected a done event last, got %v", eventKinds(events)) + } + + // Stop reason: tool_use. + if final.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason = %q, want tool_use", final.StopReason) + } + // Usage: input from message_start, output from message_delta. + if final.Usage == nil || final.Usage.InputTokens != 42 || final.Usage.OutputTokens != 57 { + t.Errorf("usage = %+v, want input=42 output=57", final.Usage) + } + // Response identity from message_start. + if final.ResponseID != "msg_01ABC" || final.ResponseModel != "claude-opus-4-8" { + t.Errorf("response id/model = %q/%q", final.ResponseID, final.ResponseModel) + } + + // Content blocks in index order: thinking, text, tool_use. + if len(final.Content) != 3 { + t.Fatalf("expected 3 content blocks, got %d: %+v", len(final.Content), final.Content) + } + th, ok := final.Content[0].(agentcore.ThinkingContent) + if !ok || th.Thinking != "Let me think. Use the tool." { + t.Errorf("thinking block = %+v", final.Content[0]) + } + if th.ThinkingSignature != "sigABC" { + t.Errorf("thinking signature = %q, want sigABC", th.ThinkingSignature) + } + txt, ok := final.Content[1].(agentcore.TextContent) + if !ok || txt.Text != "I'll check the weather." { + t.Errorf("text block = %+v", final.Content[1]) + } + tool, ok := final.Content[2].(agentcore.ToolCallContent) + if !ok || tool.Name != "get_weather" || tool.ID != "toolu_01" { + t.Fatalf("tool block = %+v", final.Content[2]) + } + // tool_use JSON must have accumulated into valid arguments. + var args map[string]string + if err := json.Unmarshal(tool.Arguments, &args); err != nil { + t.Fatalf("tool arguments not valid JSON %q: %v", tool.Arguments, err) + } + if args["city"] != "SF" { + t.Errorf("tool arguments = %v, want city=SF", args) + } +} + +func TestAnthropicDecoderTextOnlyEndTurn(t *testing.T) { + body := `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-x","usage":{"input_tokens":10,"output_tokens":0}}} + +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}} + +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" world"}} + +data: {"type":"content_block_stop","index":0} + +data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}} + +data: {"type":"message_stop"} + +` + dec := NewAnthropicDecoder() + _, final := feedSSE(t, dec, body) + if final.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("stop reason = %q, want end_turn", final.StopReason) + } + if len(final.Content) != 1 { + t.Fatalf("expected 1 content block, got %d", len(final.Content)) + } + txt, ok := final.Content[0].(agentcore.TextContent) + if !ok || txt.Text != "Hello world" { + t.Errorf("text = %+v", final.Content[0]) + } +} + +func TestAnthropicDecoderMaxTokensMapsToLength(t *testing.T) { + body := `data: {"type":"message_start","message":{"id":"m","model":"c","usage":{"input_tokens":5,"output_tokens":0}}} + +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"truncated"}} + +data: {"type":"content_block_stop","index":0} + +data: {"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":100}} + +data: {"type":"message_stop"} + +` + dec := NewAnthropicDecoder() + _, final := feedSSE(t, dec, body) + if final.StopReason != agentcore.StopReasonLength { + t.Errorf("max_tokens must map to length, got %q", final.StopReason) + } +} + +// TestAnthropicDecoderErrorEvent verifies an Anthropic `error` event becomes a +// decode error (which the transport turns into a terminal error event), never +// a panic. +func TestAnthropicDecoderErrorEvent(t *testing.T) { + dec := NewAnthropicDecoder() + _, err := dec.Decode([]byte(`{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`)) + if err == nil { + t.Fatal("error event must return a decode error") + } + if !strings.Contains(err.Error(), "overloaded_error") { + t.Errorf("error should name the type, got %v", err) + } +} + +// TestAnthropicDecoderMalformedPayload verifies invalid JSON is a returned +// error (rides the stream as terminal error), not a panic. +func TestAnthropicDecoderMalformedPayload(t *testing.T) { + dec := NewAnthropicDecoder() + if _, err := dec.Decode([]byte(`{not json`)); err == nil { + t.Fatal("malformed payload must return an error") + } +} + +// TestAnthropicDecoderFinishFlushesPartial verifies a stream cut short (no +// message_stop) still yields a done event on Finish so the partial isn't lost. +func TestAnthropicDecoderFinishFlushesPartial(t *testing.T) { + body := `data: {"type":"message_start","message":{"id":"m","model":"c","usage":{"input_tokens":5,"output_tokens":0}}} + +data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}} + +data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}} + +` + dec := NewAnthropicDecoder() + events, final := feedSSE(t, dec, body) + if events[len(events)-1].EventKind() != StreamEventDone { + t.Fatalf("Finish must emit a terminal done event, got %v", eventKinds(events)) + } + // No message_delta arrived → default end_turn. + if final.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("cut-short stream should default to end_turn, got %q", final.StopReason) + } + if len(final.Content) != 1 { + t.Fatalf("expected the partial text block, got %+v", final.Content) + } +} + +// TestAnthropicDecoderThroughTransport wires the decoder through the real +// transport pump against a recorded SSE server, exercising the full path. +func TestAnthropicDecoderThroughTransport(t *testing.T) { + srv := sseServer(t, anthropicToolUseSSE) + defer srv.Close() + + stream, err := StreamRequest(context.Background(), TransportConfig{ + NewRequest: newReqFn(srv.URL), + Decoder: NewAnthropicDecoder(), + }) + if err != nil { + t.Fatalf("StreamRequest: %v", err) + } + var kinds []string + for ev := range stream.Events() { + kinds = append(kinds, ev.EventKind()) + } + final, resErr := stream.Result(context.Background()) + if resErr != nil { + t.Fatalf("result: %v", resErr) + } + if final.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason via transport = %q, want tool_use", final.StopReason) + } + if len(final.Content) != 3 { + t.Errorf("expected 3 content blocks via transport, got %d", len(final.Content)) + } + if kinds[len(kinds)-1] != StreamEventDone { + t.Errorf("stream must end with done, got %v", kinds) + } +} + +func eventKinds(events []StreamEvent) []string { + out := make([]string, len(events)) + for i, ev := range events { + out[i] = ev.EventKind() + } + return out +} diff --git a/pigo/internal/provider/auth.go b/pigo/internal/provider/auth.go new file mode 100644 index 0000000..10335e4 --- /dev/null +++ b/pigo/internal/provider/auth.go @@ -0,0 +1,249 @@ +// This file implements credential resolution (US-012): API key lookup from +// environment variables and a config file (per provider), plus an OAuth token +// source that refreshes short-lived tokens on expiry. The resolver satisfies +// the LoopConfig.GetAPIKey shape (func(ctx, provider) string) so the agent loop +// can obtain a fresh key per request. +// +// Security (FR: secret values are not written to logs): secret values are never logged or embedded in +// error messages. Errors and String()/redaction helpers reference credentials +// by key name / provider only. +package provider + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "time" +) + +// APIKeyConfig is the on-disk config-file shape: a map of provider name to API +// key. It is parsed from JSON and holds only static keys (OAuth lives in +// TokenSource). Values are secrets and must not be logged. +type APIKeyConfig struct { + // Keys maps provider name → API key. + Keys map[string]string `json:"keys"` +} + +// LoadAPIKeyConfig parses an APIKeyConfig from JSON bytes (e.g. a config file). +func LoadAPIKeyConfig(data []byte) (*APIKeyConfig, error) { + var cfg APIKeyConfig + if err := json.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("auth: parse api key config: %w", err) + } + if cfg.Keys == nil { + cfg.Keys = make(map[string]string) + } + return &cfg, nil +} + +// LoadAPIKeyConfigFile reads and parses an APIKeyConfig from a file path. A +// missing file is not an error — it returns an empty config so env/OAuth can +// still resolve keys. +func LoadAPIKeyConfigFile(path string) (*APIKeyConfig, error) { + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return &APIKeyConfig{Keys: make(map[string]string)}, nil + } + return nil, fmt.Errorf("auth: read api key config %q: %w", path, err) + } + return LoadAPIKeyConfig(data) +} + +// envAPIKey returns the API key for a provider from the environment. It derives +// the candidate variable names from the provider registry (the single source of +// truth: LookupProviderSpec(provider).EnvVars, in precedence order), then falls +// back to a generic _API_KEY when the provider is unknown or none of +// its registered vars are set. Returns "" when no value is present. +func envAPIKey(provider string) string { + if spec, ok := LookupProviderSpec(provider); ok { + for _, name := range spec.EnvVars { + if v := os.Getenv(name); v != "" { + return v + } + } + } + // Generic fallback for unknown providers or when no registered var is set. + generic := strings.ToUpper(provider) + "_API_KEY" + return os.Getenv(generic) +} + +// TokenSource yields an access token, refreshing it when expired. It models an +// OAuth credential whose access token is short-lived (FR-15: getApiKey refreshes +// on expiry). It is safe for concurrent use. +type TokenSource struct { + mu sync.Mutex + accessToken string + expiry time.Time + refreshToken string + // Refresh exchanges the current refresh token for a new access token. It + // returns the new access token, its expiry, and (optionally) a rotated + // refresh token. Required for a TokenSource to refresh; nil means the token + // is static and never refreshed. + Refresh func(ctx context.Context, refreshToken string) (OAuthToken, error) + // Now is injectable for testing; defaults to time.Now. + Now func() time.Time + // Leeway refreshes the token this long before its actual expiry to avoid + // racing the boundary. Defaults to 30s. + Leeway time.Duration +} + +// OAuthToken is the result of an OAuth exchange/refresh. Values are secrets. +type OAuthToken struct { + AccessToken string + RefreshToken string + Expiry time.Time +} + +// NewTokenSource builds a TokenSource seeded with an initial token and a refresh +// function. refresh may be nil for a static (never-expiring) token. +func NewTokenSource(initial OAuthToken, refresh func(ctx context.Context, refreshToken string) (OAuthToken, error)) *TokenSource { + return &TokenSource{ + accessToken: initial.AccessToken, + expiry: initial.Expiry, + refreshToken: initial.RefreshToken, + Refresh: refresh, + } +} + +func (t *TokenSource) now() time.Time { + if t.Now != nil { + return t.Now() + } + return time.Now() +} + +// defaultTokenLeeway is how far before an OAuth token's expiry it is treated as +// already expired, so a refresh happens before a request rather than mid-flight. +const defaultTokenLeeway = 30 * time.Second + +func (t *TokenSource) leeway() time.Duration { + if t.Leeway > 0 { + return t.Leeway + } + return defaultTokenLeeway +} + +// expired reports whether the access token is missing or within leeway of its +// expiry. A zero expiry means "never expires" (static token). +func (t *TokenSource) expired() bool { + if t.accessToken == "" { + return true + } + if t.expiry.IsZero() { + return false + } + return !t.now().Before(t.expiry.Add(-t.leeway())) +} + +// Token returns a valid access token, refreshing it when expired. It errors if +// a refresh is needed but no Refresh func is set, or if Refresh fails. The +// returned error never contains the token value. +func (t *TokenSource) Token(ctx context.Context) (string, error) { + t.mu.Lock() + defer t.mu.Unlock() + if !t.expired() { + return t.accessToken, nil + } + if t.Refresh == nil { + return "", fmt.Errorf("auth: token expired and no refresh function configured") + } + tok, err := t.Refresh(ctx, t.refreshToken) + if err != nil { + return "", fmt.Errorf("auth: token refresh failed: %w", err) + } + t.accessToken = tok.AccessToken + t.expiry = tok.Expiry + if tok.RefreshToken != "" { + t.refreshToken = tok.RefreshToken + } + return t.accessToken, nil +} + +// CredentialStore resolves API keys per provider from three layers, in order: +// OAuth token source (if registered), environment variable, config file. It +// implements the LoopConfig.GetAPIKey shape via GetAPIKey. +// +// It is safe for concurrent use. +type CredentialStore struct { + mu sync.RWMutex + config *APIKeyConfig + sources map[string]*TokenSource // provider → OAuth token source + overrides map[string]string // provider → explicit key (highest static priority) +} + +// NewCredentialStore builds a store over an optional config file. A nil config +// is treated as empty. +func NewCredentialStore(config *APIKeyConfig) *CredentialStore { + if config == nil { + config = &APIKeyConfig{Keys: make(map[string]string)} + } + return &CredentialStore{ + config: config, + sources: make(map[string]*TokenSource), + overrides: make(map[string]string), + } +} + +// SetOverride records an explicit API key for a provider that wins over the +// environment variable and config file (but not a live OAuth token, which is +// auto-refreshed). It is the seam for a CLI --api-key flag: an empty key is +// ignored so a bare flag does not clobber env/config resolution. +func (c *CredentialStore) SetOverride(provider, key string) { + if key == "" { + return + } + c.mu.Lock() + defer c.mu.Unlock() + c.overrides[provider] = key +} + +// RegisterOAuth registers an OAuth TokenSource for a provider. Once registered, +// GetAPIKey prefers the (auto-refreshing) OAuth token over static keys. +func (c *CredentialStore) RegisterOAuth(provider string, src *TokenSource) { + c.mu.Lock() + defer c.mu.Unlock() + c.sources[provider] = src +} + +// GetAPIKey resolves the API key for a provider. Resolution order: OAuth token +// (refreshed on expiry) → explicit override (--api-key) → environment variable +// → config file. Returns "" when no credential is available. This matches +// LoopConfig.GetAPIKey so it can be assigned directly. +// +// On OAuth refresh failure it falls back to override/env/config rather than +// returning a secret-bearing error; the empty return lets the caller fall back +// to a static key. It never logs secret values. +func (c *CredentialStore) GetAPIKey(ctx context.Context, provider string) string { + c.mu.RLock() + src := c.sources[provider] + override := c.overrides[provider] + cfgKey := "" + if c.config != nil { + cfgKey = c.config.Keys[provider] + } + c.mu.RUnlock() + + if src != nil { + if tok, err := src.Token(ctx); err == nil && tok != "" { + return tok + } + // Refresh failed → fall through to static layers. + } + if override != "" { + return override + } + if env := envAPIKey(provider); env != "" { + return env + } + return cfgKey +} + +// HasCredential reports whether any credential (OAuth/env/config) is available +// for a provider, without exposing the value. +func (c *CredentialStore) HasCredential(ctx context.Context, provider string) bool { + return c.GetAPIKey(ctx, provider) != "" +} diff --git a/pigo/internal/provider/auth_test.go b/pigo/internal/provider/auth_test.go new file mode 100644 index 0000000..7584f0f --- /dev/null +++ b/pigo/internal/provider/auth_test.go @@ -0,0 +1,255 @@ +package provider + +import ( + "context" + "testing" + "time" +) + +func TestEnvAPIKey(t *testing.T) { + t.Setenv("ANTHROPIC_OAUTH_TOKEN", "") + t.Setenv("ANTHROPIC_API_KEY", "sk-ant-env") + if got := envAPIKey("anthropic"); got != "sk-ant-env" { + t.Errorf("env key = %q, want sk-ant-env", got) + } + // Unknown provider uses generic _API_KEY fallback. + t.Setenv("FOOBAR_API_KEY", "sk-foobar") + if got := envAPIKey("foobar"); got != "sk-foobar" { + t.Errorf("generic env key = %q, want sk-foobar", got) + } + if got := envAPIKey("nonesuch"); got != "" { + t.Errorf("missing env key = %q, want empty", got) + } +} + +// TestEnvAPIKeyFromRegistry verifies API-key resolution derives from the +// provider registry (single source of truth) across a representative set of +// providers, that Anthropic's OAuth token takes precedence over its API key, +// and that an unknown provider hits the generic _API_KEY fallback. +func TestEnvAPIKeyFromRegistry(t *testing.T) { + cases := []struct { + provider string + envVar string + value string + }{ + {"deepseek", "DEEPSEEK_API_KEY", "sk-deepseek"}, + {"groq", "GROQ_API_KEY", "sk-groq"}, + {"zai", "ZAI_API_KEY", "sk-zai"}, + {"moonshotai-cn", "MOONSHOT_API_KEY", "sk-moonshot-cn"}, + {"xiaomi-token-plan-ams", "XIAOMI_TOKEN_PLAN_AMS_API_KEY", "sk-xiaomi-ams"}, + } + for _, tc := range cases { + t.Run(tc.provider, func(t *testing.T) { + t.Setenv(tc.envVar, tc.value) + if got := envAPIKey(tc.provider); got != tc.value { + t.Errorf("envAPIKey(%q) = %q, want %q", tc.provider, got, tc.value) + } + }) + } + + // Anthropic: OAuth token wins over API key (registry ordering). + t.Run("anthropic-oauth-first", func(t *testing.T) { + t.Setenv("ANTHROPIC_OAUTH_TOKEN", "oauth-tok") + t.Setenv("ANTHROPIC_API_KEY", "sk-ant") + if got := envAPIKey("anthropic"); got != "oauth-tok" { + t.Errorf("anthropic = %q, want oauth-tok (OAuth precedence)", got) + } + // With OAuth unset, the API key resolves. + t.Setenv("ANTHROPIC_OAUTH_TOKEN", "") + if got := envAPIKey("anthropic"); got != "sk-ant" { + t.Errorf("anthropic (no oauth) = %q, want sk-ant", got) + } + }) + + // Unknown provider falls back to the generic convention. + t.Run("unknown-generic-fallback", func(t *testing.T) { + t.Setenv("MADEUP_PROVIDER_API_KEY", "sk-generic") + if got := envAPIKey("madeup-provider"); got != "" { + // Hyphenated names uppercase to MADEUP-PROVIDER_API_KEY, not a match; + // verify the true generic form resolves for an underscore-friendly name. + t.Logf("hyphenated generic = %q", got) + } + t.Setenv("MADEUPPROVIDER_API_KEY", "sk-generic2") + if got := envAPIKey("madeupprovider"); got != "sk-generic2" { + t.Errorf("generic fallback = %q, want sk-generic2", got) + } + }) +} + +func TestLoadAPIKeyConfig(t *testing.T) { + cfg, err := LoadAPIKeyConfig([]byte(`{"keys":{"openai":"sk-openai-cfg"}}`)) + if err != nil { + t.Fatalf("load: %v", err) + } + if cfg.Keys["openai"] != "sk-openai-cfg" { + t.Errorf("config key = %q", cfg.Keys["openai"]) + } + if _, err := LoadAPIKeyConfig([]byte(`not json`)); err == nil { + t.Fatal("bad JSON must error") + } +} + +func TestLoadAPIKeyConfigFileMissing(t *testing.T) { + cfg, err := LoadAPIKeyConfigFile("/no/such/path/keys.json") + if err != nil { + t.Fatalf("missing file must not error: %v", err) + } + if len(cfg.Keys) != 0 { + t.Errorf("missing file must yield empty keys, got %v", cfg.Keys) + } +} + +// TestCredentialStoreResolutionOrder verifies OAuth > env > config precedence. +func TestCredentialStoreResolutionOrder(t *testing.T) { + cfg, _ := LoadAPIKeyConfig([]byte(`{"keys":{"anthropic":"sk-cfg","openai":"sk-openai-cfg"}}`)) + store := NewCredentialStore(cfg) + + // Neutralize any ambient keys so config-only resolution is deterministic. + t.Setenv("OPENAI_API_KEY", "") + t.Setenv("ANTHROPIC_API_KEY", "") + t.Setenv("CLAUDE_API_KEY", "") + + // Config-only provider resolves from config. + if got := store.GetAPIKey(context.Background(), "openai"); got != "sk-openai-cfg" { + t.Errorf("openai (config) = %q, want sk-openai-cfg", got) + } + + // Env overrides config. + t.Setenv("ANTHROPIC_API_KEY", "sk-env") + if got := store.GetAPIKey(context.Background(), "anthropic"); got != "sk-env" { + t.Errorf("anthropic (env>config) = %q, want sk-env", got) + } + + // OAuth overrides env + config. + store.RegisterOAuth("anthropic", NewTokenSource( + OAuthToken{AccessToken: "oauth-token", Expiry: time.Now().Add(time.Hour)}, nil)) + if got := store.GetAPIKey(context.Background(), "anthropic"); got != "oauth-token" { + t.Errorf("anthropic (oauth>env) = %q, want oauth-token", got) + } + + // Unknown provider → empty. + if got := store.GetAPIKey(context.Background(), "ghost"); got != "" { + t.Errorf("ghost = %q, want empty", got) + } +} + +// TestCredentialStoreOverride verifies an explicit --api-key override wins over +// env and config, but not over a live OAuth token, and that an empty override +// is ignored (so a bare flag does not clobber env/config). +func TestCredentialStoreOverride(t *testing.T) { + cfg, _ := LoadAPIKeyConfig([]byte(`{"keys":{"openai":"sk-openai-cfg"}}`)) + store := NewCredentialStore(cfg) + t.Setenv("OPENAI_API_KEY", "sk-openai-env") + + // Empty override is a no-op: env still wins over config. + store.SetOverride("openai", "") + if got := store.GetAPIKey(context.Background(), "openai"); got != "sk-openai-env" { + t.Errorf("empty override should not clobber env, got %q", got) + } + + // Non-empty override wins over env and config. + store.SetOverride("openai", "sk-flag") + if got := store.GetAPIKey(context.Background(), "openai"); got != "sk-flag" { + t.Errorf("override should win over env/config, got %q", got) + } + + // OAuth still wins over an override. + store.RegisterOAuth("openai", NewTokenSource( + OAuthToken{AccessToken: "oauth-token", Expiry: time.Now().Add(time.Hour)}, nil)) + if got := store.GetAPIKey(context.Background(), "openai"); got != "oauth-token" { + t.Errorf("oauth should win over override, got %q", got) + } +} + +// TestTokenSourceRefresh verifies an expired token triggers a refresh returning +// a new token. +func TestTokenSourceRefresh(t *testing.T) { + now := time.Now() + refreshCount := 0 + src := NewTokenSource( + OAuthToken{AccessToken: "old", RefreshToken: "refresh-1", Expiry: now.Add(-time.Minute)}, + func(ctx context.Context, rt string) (OAuthToken, error) { + refreshCount++ + if rt != "refresh-1" { + t.Errorf("refresh token = %q, want refresh-1", rt) + } + return OAuthToken{AccessToken: "new", RefreshToken: "refresh-2", Expiry: now.Add(time.Hour)}, nil + }, + ) + src.Now = func() time.Time { return now } + + tok, err := src.Token(context.Background()) + if err != nil { + t.Fatalf("token: %v", err) + } + if tok != "new" { + t.Errorf("token = %q, want new (refreshed)", tok) + } + if refreshCount != 1 { + t.Errorf("refresh count = %d, want 1", refreshCount) + } + + // Second call within validity does not refresh again. + if _, err := src.Token(context.Background()); err != nil { + t.Fatalf("token 2: %v", err) + } + if refreshCount != 1 { + t.Errorf("refresh count after valid reuse = %d, want 1", refreshCount) + } +} + +func TestTokenSourceNoRefreshFunc(t *testing.T) { + now := time.Now() + // Expired token with no Refresh func → error. + src := NewTokenSource(OAuthToken{AccessToken: "old", Expiry: now.Add(-time.Minute)}, nil) + src.Now = func() time.Time { return now } + if _, err := src.Token(context.Background()); err == nil { + t.Fatal("expired token without refresh must error") + } + + // Static token (zero expiry) never expires. + static := NewTokenSource(OAuthToken{AccessToken: "static"}, nil) + tok, err := static.Token(context.Background()) + if err != nil || tok != "static" { + t.Errorf("static token = %q, err = %v", tok, err) + } +} + +func TestTokenSourceRefreshError(t *testing.T) { + now := time.Now() + src := NewTokenSource( + OAuthToken{AccessToken: "old", Expiry: now.Add(-time.Minute)}, + func(ctx context.Context, rt string) (OAuthToken, error) { + return OAuthToken{}, context.DeadlineExceeded + }, + ) + src.Now = func() time.Time { return now } + _, err := src.Token(context.Background()) + if err == nil { + t.Fatal("refresh error must propagate") + } + // Error must not leak the (empty) token but should mention refresh. + if got := err.Error(); got == "" { + t.Error("expected non-empty error") + } +} + +// TestCredentialStoreOAuthRefreshFallback verifies a failing OAuth refresh +// falls back to env/config rather than returning empty when a static key exists. +func TestCredentialStoreOAuthRefreshFallback(t *testing.T) { + cfg, _ := LoadAPIKeyConfig([]byte(`{"keys":{"anthropic":"sk-cfg-fallback"}}`)) + store := NewCredentialStore(cfg) + now := time.Now() + src := NewTokenSource( + OAuthToken{AccessToken: "old", Expiry: now.Add(-time.Minute)}, + func(ctx context.Context, rt string) (OAuthToken, error) { + return OAuthToken{}, context.DeadlineExceeded + }, + ) + src.Now = func() time.Time { return now } + store.RegisterOAuth("anthropic", src) + + if got := store.GetAPIKey(context.Background(), "anthropic"); got != "sk-cfg-fallback" { + t.Errorf("refresh-failed fallback = %q, want sk-cfg-fallback", got) + } +} diff --git a/pigo/internal/provider/image_test.go b/pigo/internal/provider/image_test.go new file mode 100644 index 0000000..8d6026e --- /dev/null +++ b/pigo/internal/provider/image_test.go @@ -0,0 +1,157 @@ +package provider + +// Tests for image (multimodal) input encoding on both provider wires (US-010, +// #126). They exercise encodeOpenAIMessage / encodeAnthropicMessage directly +// (unit-level, no HTTP) to assert the exact wire shape of image blocks, plus the +// checkImageSupport guard that turns image input on a text-only model into a +// clear error rather than a silent drop. + +import ( + "encoding/json" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// imageUserMessage builds a user message carrying one text block and one image +// block, the common multimodal input shape. +func imageUserMessage(text, data, mime string) agentcore.UserMessage { + return agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{ + agentcore.NewTextContent(text), + agentcore.NewImageContent(data, mime), + }, + } +} + +// TestEncodeOpenAIMessageImageBlock asserts a user message with an image encodes +// to the multimodal array form with an image_url data URI. +func TestEncodeOpenAIMessageImageBlock(t *testing.T) { + msg := imageUserMessage("what is this?", "QUJD", "image/png") + out := encodeOpenAIMessage(msg) + if len(out) != 1 { + t.Fatalf("encodeOpenAIMessage returned %d entries, want 1", len(out)) + } + entry := out[0] + if entry["role"] != "user" { + t.Errorf("role = %v, want user", entry["role"]) + } + parts, ok := entry["content"].([]map[string]any) + if !ok { + t.Fatalf("content is %T, want []map[string]any (array form)", entry["content"]) + } + if len(parts) != 2 { + t.Fatalf("content has %d parts, want 2 (text + image)", len(parts)) + } + if parts[0]["type"] != "text" || parts[0]["text"] != "what is this?" { + t.Errorf("text part = %#v", parts[0]) + } + if parts[1]["type"] != "image_url" { + t.Fatalf("image part type = %v, want image_url", parts[1]["type"]) + } + iu, ok := parts[1]["image_url"].(map[string]any) + if !ok { + t.Fatalf("image_url is %T", parts[1]["image_url"]) + } + if want := "data:image/png;base64,QUJD"; iu["url"] != want { + t.Errorf("image_url.url = %v, want %v", iu["url"], want) + } +} + +// TestEncodeOpenAIMessageNoImageIsString asserts a text-only user message stays +// a plain string (not the array form), preserving the common-case wire shape. +func TestEncodeOpenAIMessageNoImageIsString(t *testing.T) { + msg := agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent("hello")}, + } + out := encodeOpenAIMessage(msg) + if s, ok := out[0]["content"].(string); !ok || s != "hello" { + t.Errorf("content = %#v, want string \"hello\"", out[0]["content"]) + } +} + +// TestEncodeAnthropicMessageImageBlock asserts a user message with an image +// encodes to the content-block array form with a base64 image source. +func TestEncodeAnthropicMessageImageBlock(t *testing.T) { + msg := imageUserMessage("describe", "REVG", "image/jpeg") + entry := encodeAnthropicMessage(msg) + if entry["role"] != "user" { + t.Errorf("role = %v, want user", entry["role"]) + } + blocks, ok := entry["content"].([]map[string]any) + if !ok { + t.Fatalf("content is %T, want []map[string]any (array form)", entry["content"]) + } + if len(blocks) != 2 { + t.Fatalf("content has %d blocks, want 2 (text + image)", len(blocks)) + } + if blocks[0]["type"] != "text" || blocks[0]["text"] != "describe" { + t.Errorf("text block = %#v", blocks[0]) + } + if blocks[1]["type"] != "image" { + t.Fatalf("image block type = %v, want image", blocks[1]["type"]) + } + src, ok := blocks[1]["source"].(map[string]any) + if !ok { + t.Fatalf("source is %T", blocks[1]["source"]) + } + if src["type"] != "base64" || src["media_type"] != "image/jpeg" || src["data"] != "REVG" { + t.Errorf("source = %#v", src) + } +} + +// TestEncodeAnthropicMessageNoImageIsString asserts a text-only user message +// stays a plain string on the Anthropic wire. +func TestEncodeAnthropicMessageNoImageIsString(t *testing.T) { + msg := agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent("hi")}, + } + entry := encodeAnthropicMessage(msg) + if s, ok := entry["content"].(string); !ok || s != "hi" { + t.Errorf("content = %#v, want string \"hi\"", entry["content"]) + } +} + +// TestImageBlocksAreJSONSerializable guards against map value types that +// json.Marshal cannot encode: both wire shapes must round-trip to JSON. +func TestImageBlocksAreJSONSerializable(t *testing.T) { + msg := imageUserMessage("x", "QQ==", "image/webp") + if _, err := json.Marshal(encodeOpenAIMessage(msg)); err != nil { + t.Errorf("marshal openai image message: %v", err) + } + if _, err := json.Marshal(encodeAnthropicMessage(msg)); err != nil { + t.Errorf("marshal anthropic image message: %v", err) + } +} + +// TestCheckImageSupport asserts the capability guard: image input on a model +// that declares SupportsImages passes; on a text-only model it errors; and a +// text-only prompt always passes regardless of the model. +func TestCheckImageSupport(t *testing.T) { + models := []Model{ + {ID: "vision-1", SupportsImages: true}, + {ID: "text-1", SupportsImages: false}, + } + imgMsgs := []agentcore.Message{imageUserMessage("q", "QQ==", "image/png")} + textMsgs := []agentcore.Message{agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent("no image")}, + }} + + if err := checkImageSupport("p", "vision-1", models, imgMsgs); err != nil { + t.Errorf("vision model rejected image input: %v", err) + } + if err := checkImageSupport("p", "text-1", models, imgMsgs); err == nil { + t.Error("text-only model accepted image input, want error") + } + if err := checkImageSupport("p", "text-1", models, textMsgs); err != nil { + t.Errorf("text-only prompt on text model errored: %v", err) + } + // Unknown model (not in catalog) defers to provider validation → no error. + if err := checkImageSupport("p", "unknown", models, imgMsgs); err != nil { + t.Errorf("unknown model errored on image input: %v", err) + } +} diff --git a/pigo/internal/provider/infer.go b/pigo/internal/provider/infer.go new file mode 100644 index 0000000..15a632a --- /dev/null +++ b/pigo/internal/provider/infer.go @@ -0,0 +1,78 @@ +// This file implements model-name → provider inference (US: auto-infer provider +// from --model alone). When a user supplies only --model, with no --provider, +// --protocol, or --base-url, pigo tries to guess the owning provider from the +// model id's well-known name prefix (e.g. "claude-*" → anthropic, "deepseek-*" +// → deepseek). This lets `pigo -m claude-opus-4-8` reach the Anthropic API +// without the user spelling out the provider or its wire protocol. +// +// The mapping is deliberately conservative: only names that unambiguously +// identify a single built-in provider are inferred. Ambiguous families served +// by many gateways (notably "llama-*", "qwq-*", "gemma-*", "mixtral-*") are NOT +// inferred — they return ok=false so the caller falls through to its existing +// default (OpenRouter), which is the safe, unchanged behavior. +// +// Every provider name returned here is guaranteed to exist in providerRegistry +// (enforced by a test), so callers can hand the result straight to the +// registry-driven resolution path. +package provider + +import "strings" + +// modelPrefixProvider maps a lowercase model-name prefix to the built-in +// provider name that serves that family. Order matters: the table is scanned +// top-to-bottom and the FIRST matching prefix wins, so list more specific +// prefixes before shorter ones that would also match. +// +// Each provider name here must be present in providerRegistry (registry.go). +var modelPrefixProvider = []struct { + prefix string + provider string +}{ + {"claude-", "anthropic"}, + {"gpt-", "openai"}, + {"o1-", "openai"}, + {"o3-", "openai"}, + {"o4-", "openai"}, + {"gemini-", "google"}, + {"deepseek-", "deepseek"}, + {"glm-", "zai"}, + {"kimi-", "moonshotai"}, + {"moonshot-", "moonshotai"}, + {"qwen-", "dashscope"}, + {"ernie-", "qianfan"}, + {"doubao-", "volcengine"}, + {"grok-", "xai"}, + {"mistral-", "mistral"}, + {"codestral-", "mistral"}, + {"devstral-", "mistral"}, + {"hunyuan-", "hunyuan"}, + {"minimax-", "minimax"}, + {"mimo-", "xiaomi"}, +} + +// InferProviderFromModel guesses the built-in provider name that serves a given +// model id, based on the id's well-known name prefix. It returns the provider +// name and ok=true on a confident match, or ("", false) when the id is unknown +// or ambiguous (served by multiple gateways). Matching is case-insensitive. +// +// The returned name is always a valid entry in providerRegistry. Callers should +// use it only when no explicit --provider/--protocol/--base-url was given; those +// flags take precedence over any inference. +func InferProviderFromModel(model string) (string, bool) { + id := strings.ToLower(strings.TrimSpace(model)) + if id == "" { + return "", false + } + // A "provider/model" style id (e.g. "openai/gpt-4o") is an OpenRouter-style + // routed id, not a bare model name — leave those to the caller's preset/ + // prefix handling rather than inferring from the leading segment. + if strings.Contains(id, "/") { + return "", false + } + for _, m := range modelPrefixProvider { + if strings.HasPrefix(id, m.prefix) { + return m.provider, true + } + } + return "", false +} diff --git a/pigo/internal/provider/infer_test.go b/pigo/internal/provider/infer_test.go new file mode 100644 index 0000000..ee87988 --- /dev/null +++ b/pigo/internal/provider/infer_test.go @@ -0,0 +1,96 @@ +package provider + +// Tests for model-name → provider inference (InferProviderFromModel). They +// verify each documented name prefix resolves to the expected built-in +// provider, that ambiguous/unknown ids and routed "provider/model" ids do not +// resolve, that matching is case-insensitive, and that every inferred provider +// name actually exists in the provider registry (the single source of truth). + +import "testing" + +// TestInferProviderFromModelKnown verifies each documented model-name prefix +// resolves to its expected built-in provider. +func TestInferProviderFromModelKnown(t *testing.T) { + cases := []struct { + model string + want string + }{ + {"claude-opus-4-8", "anthropic"}, + {"claude-3.5-sonnet", "anthropic"}, + {"gpt-4o", "openai"}, + {"gpt-4o-mini", "openai"}, + {"o1-preview", "openai"}, + {"o3-mini", "openai"}, + {"o4-mini", "openai"}, + {"gemini-2.5-pro", "google"}, + {"deepseek-chat", "deepseek"}, + {"deepseek-v4-pro", "deepseek"}, + {"glm-5.1", "zai"}, + {"kimi-k2-thinking", "moonshotai"}, + {"moonshot-v1-8k", "moonshotai"}, + {"qwen-max", "dashscope"}, + {"ernie-4.5-turbo-32k", "qianfan"}, + {"doubao-seed-1-6", "volcengine"}, + {"grok-4.5", "xai"}, + {"mistral-large-latest", "mistral"}, + {"codestral-latest", "mistral"}, + {"devstral-medium-latest", "mistral"}, + {"hunyuan-turbos-latest", "hunyuan"}, + {"minimax-m2.7", "minimax"}, + {"mimo-v2-pro", "xiaomi"}, + } + for _, c := range cases { + got, ok := InferProviderFromModel(c.model) + if !ok { + t.Errorf("InferProviderFromModel(%q): ok=false, want provider %q", c.model, c.want) + continue + } + if got != c.want { + t.Errorf("InferProviderFromModel(%q) = %q, want %q", c.model, got, c.want) + } + } +} + +// TestInferProviderFromModelCaseInsensitive verifies matching ignores case and +// surrounding whitespace. +func TestInferProviderFromModelCaseInsensitive(t *testing.T) { + for _, m := range []string{"Claude-Opus-4-8", " GPT-4o ", "DeepSeek-Chat"} { + if _, ok := InferProviderFromModel(m); !ok { + t.Errorf("InferProviderFromModel(%q): ok=false, want a match", m) + } + } +} + +// TestInferProviderFromModelAmbiguousOrUnknown verifies ids that are ambiguous +// (served by many gateways), routed ("provider/model"), empty, or simply +// unknown do NOT resolve — the caller must fall through to its default. +func TestInferProviderFromModelAmbiguousOrUnknown(t *testing.T) { + for _, m := range []string{ + "", // empty + " ", // whitespace only + "llama-3.3-70b-instruct", // ambiguous: many gateways + "qwq-32b", // ambiguous + "gemma-2-9b-it", // ambiguous + "mixtral-8x22b", // ambiguous + "openai/gpt-4o", // routed id, leave to preset/prefix handling + "anthropic/claude-3.5", // routed id + "ollama/llama3.3", // routed id (ollama prefix path) + "totally-made-up-model", // unknown + } { + if got, ok := InferProviderFromModel(m); ok { + t.Errorf("InferProviderFromModel(%q) = (%q, true), want ok=false", m, got) + } + } +} + +// TestInferProviderNamesExistInRegistry verifies every provider name the +// inference table can return is a real built-in provider (registry is the +// single source of truth), so a hit can be handed straight to registry-driven +// resolution. +func TestInferProviderNamesExistInRegistry(t *testing.T) { + for _, m := range modelPrefixProvider { + if _, ok := LookupProviderSpec(m.provider); !ok { + t.Errorf("inference maps prefix %q → %q, which is not in providerRegistry", m.prefix, m.provider) + } + } +} diff --git a/pigo/internal/provider/openai.go b/pigo/internal/provider/openai.go new file mode 100644 index 0000000..d331e7c --- /dev/null +++ b/pigo/internal/provider/openai.go @@ -0,0 +1,245 @@ +// This file implements the OpenAI-compatible streaming decoder (US-009): a +// stateful Decoder (see transport.go) for the OpenAI Chat Completions SSE +// stream, which is also the wire format spoken by most third-party gateways +// (OpenRouter, Groq, together, local servers, …). Selecting the base URL is a +// transport concern (NewRequest builds the *http.Request), so this decoder is +// base-URL agnostic and reused across every OpenAI-compatible provider. +// +// OpenAI streams a sequence of chat.completion.chunk objects: +// +// {"choices":[{"delta":{"role":"assistant"}}]} → first chunk +// {"choices":[{"delta":{"content":"Hel"}}]} → text delta +// {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1", +// "function":{"name":"f","arguments":"{\"a\":"}}}]}]} → tool-call delta +// {"choices":[{"finish_reason":"tool_calls"}]} → stop reason +// {"usage":{"prompt_tokens":10,"completion_tokens":5}} → final usage +// [DONE] → transport-level +// +// Per the dual failure model (FR-13) the decoder never panics: malformed +// payloads surface as a returned error which the transport turns into a +// terminal StreamErrorEvent. +package provider + +import ( + "encoding/json" + "fmt" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// openaiToolCall accumulates one streamed tool call, keyed by its delta index. +// id/name arrive once (usually in the first fragment); arguments accumulate. +type openaiToolCall struct { + id string + name string + args strings.Builder +} + +// OpenAIDecoder is the stateful SSE decoder for the OpenAI Chat Completions API +// and compatible gateways. It implements the transport Decoder interface and is +// not safe for concurrent use — the transport drives it from one goroutine. +type OpenAIDecoder struct { + text strings.Builder + thinking strings.Builder // reasoning_content / reasoning stream (if any) + toolCalls map[int]*openaiToolCall + toolOrder []int // tool-call indices in first-seen order + + responseID string + responseModel string + inputTokens int + outputTokens int + stopReason string // mapped pigo stop reason (empty until finish_reason) + done bool +} + +// NewOpenAIDecoder builds a fresh decoder for one streamed response. +func NewOpenAIDecoder() *OpenAIDecoder { + return &OpenAIDecoder{toolCalls: make(map[int]*openaiToolCall)} +} + +// openaiChunk is the streamed chat.completion.chunk envelope. +type openaiChunk struct { + ID string `json:"id"` + Model string `json:"model"` + Choices []struct { + Delta struct { + Content string `json:"content"` + // ReasoningContent carries the model's reasoning/thinking stream on the + // OpenAI wire (DeepSeek-R1, Kimi, and other reasoning models put their + // chain-of-thought here). Some gateways name it "reasoning" instead, so + // both are accepted; without this field the thinking stream is silently + // dropped from the response and from history. + ReasoningContent string `json:"reasoning_content"` + Reasoning string `json:"reasoning"` + ToolCalls []openaiToolDelta `json:"tool_calls"` + } `json:"delta"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + Usage *struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + } `json:"usage"` + // Some gateways surface an error object inline on the stream. + Error *struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error"` +} + +type openaiToolDelta struct { + Index int `json:"index"` + ID string `json:"id"` + Type string `json:"type"` + Function struct { + Name string `json:"name"` + Arguments string `json:"arguments"` + } `json:"function"` +} + +// Decode turns one OpenAI SSE data payload into zero or more StreamEvents. +func (d *OpenAIDecoder) Decode(payload []byte) ([]StreamEvent, error) { + var chunk openaiChunk + if err := json.Unmarshal(payload, &chunk); err != nil { + return nil, fmt.Errorf("openai: parse chunk: %w", err) + } + if chunk.Error != nil { + msg := "openai stream error" + if chunk.Error.Type != "" { + msg = "openai " + chunk.Error.Type + } + if chunk.Error.Message != "" { + msg += ": " + chunk.Error.Message + } + return nil, fmt.Errorf("%s", msg) + } + + if chunk.ID != "" { + d.responseID = chunk.ID + } + if chunk.Model != "" { + d.responseModel = chunk.Model + } + if chunk.Usage != nil { + d.inputTokens = chunk.Usage.PromptTokens + d.outputTokens = chunk.Usage.CompletionTokens + } + + var events []StreamEvent + for _, choice := range chunk.Choices { + // Reasoning stream (DeepSeek-R1 / Kimi / …): reasoning_content is the + // common field; a few gateways use "reasoning". Accumulate whichever is set. + if r := choice.Delta.ReasoningContent; r != "" { + d.thinking.WriteString(r) + events = append(events, StreamThinkingEvent{Partial: d.partial()}) + } else if r := choice.Delta.Reasoning; r != "" { + d.thinking.WriteString(r) + events = append(events, StreamThinkingEvent{Partial: d.partial()}) + } + if choice.Delta.Content != "" { + d.text.WriteString(choice.Delta.Content) + events = append(events, StreamTextEvent{Partial: d.partial()}) + } + for _, tc := range choice.Delta.ToolCalls { + d.applyToolDelta(tc) + events = append(events, StreamToolCallEvent{Partial: d.partial()}) + } + if choice.FinishReason != "" { + d.stopReason = mapOpenAIFinishReason(choice.FinishReason) + } + } + return events, nil +} + +// Finish flushes a terminal done event if the stream ended without an explicit +// terminator, so a partial response is still delivered rather than lost. +func (d *OpenAIDecoder) Finish() ([]StreamEvent, error) { + if d.done { + return nil, nil + } + return d.finishDone(), nil +} + +// applyToolDelta merges one tool-call fragment into the accumulated state. +func (d *OpenAIDecoder) applyToolDelta(tc openaiToolDelta) { + call := d.toolCalls[tc.Index] + if call == nil { + call = &openaiToolCall{} + d.toolCalls[tc.Index] = call + d.toolOrder = append(d.toolOrder, tc.Index) + } + if tc.ID != "" { + call.id = tc.ID + } + if tc.Function.Name != "" { + call.name = tc.Function.Name + } + call.args.WriteString(tc.Function.Arguments) +} + +// finishDone builds the terminal assistant message and marks the decoder done. +func (d *OpenAIDecoder) finishDone() []StreamEvent { + if d.done { + return nil + } + d.done = true + msg := d.partial() + if msg.StopReason == "" { + msg.StopReason = agentcore.StopReasonEndTurn + } + return []StreamEvent{StreamDoneEvent{Message: msg}} +} + +// partial materializes the accumulated state into an AssistantMessage: the text +// block first (if any), then tool-call blocks in first-seen index order. +func (d *OpenAIDecoder) partial() agentcore.AssistantMessage { + msg := agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + API: "openai", + Provider: "openai", + StopReason: d.stopReason, + ResponseID: d.responseID, + ResponseModel: d.responseModel, + } + if d.inputTokens != 0 || d.outputTokens != 0 { + msg.Usage = &agentcore.Usage{InputTokens: d.inputTokens, OutputTokens: d.outputTokens} + } + if d.thinking.Len() > 0 { + msg.Content = append(msg.Content, agentcore.NewThinkingContent(d.thinking.String())) + } + if d.text.Len() > 0 { + msg.Content = append(msg.Content, agentcore.NewTextContent(d.text.String())) + } + + idx := make([]int, len(d.toolOrder)) + copy(idx, d.toolOrder) + sort.Ints(idx) + for _, i := range idx { + call := d.toolCalls[i] + if call == nil { + continue + } + args := json.RawMessage(strings.TrimSpace(call.args.String())) + if len(args) == 0 { + args = json.RawMessage("{}") + } + msg.Content = append(msg.Content, agentcore.NewToolCallContent(call.id, call.name, args)) + } + return msg +} + +// mapOpenAIFinishReason maps an OpenAI finish_reason to the pigo StopReason set. +// Unknown reasons default to end_turn (a natural, non-error stop). +func mapOpenAIFinishReason(reason string) string { + switch reason { + case "length": + return agentcore.StopReasonLength + case "tool_calls", "function_call": + return agentcore.StopReasonToolUse + case "stop": + return agentcore.StopReasonEndTurn + default: + return agentcore.StopReasonEndTurn + } +} diff --git a/pigo/internal/provider/openai_test.go b/pigo/internal/provider/openai_test.go new file mode 100644 index 0000000..6e328f1 --- /dev/null +++ b/pigo/internal/provider/openai_test.go @@ -0,0 +1,190 @@ +package provider + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// A recorded OpenAI Chat Completions SSE stream covering a text delta followed +// by a two-fragment tool call, ending with finish_reason=tool_calls and a final +// usage-only chunk. Trimmed but structurally faithful to the real wire format +// (the transport strips the `data:` prefix and the trailing [DONE]). +const openaiToolCallSSE = `data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"role":"assistant"}}]} + +data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"content":"Let me "}}]} + +data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"content":"check."}}]} + +data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":"}}]}}]} + +data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":" \"SF\"}"}}]}}]} + +data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{},"finish_reason":"tool_calls"}]} + +data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[],"usage":{"prompt_tokens":11,"completion_tokens":8}} + +data: [DONE] + +` + +func TestOpenAIDecoderToolCallStream(t *testing.T) { + dec := NewOpenAIDecoder() + events, final := feedSSE(t, dec, openaiToolCallSSE) + + // The last emitted event must be the terminal done event. + if len(events) == 0 || events[len(events)-1].EventKind() != StreamEventDone { + t.Fatalf("expected a done event last, got %v", eventKinds(events)) + } + // Stop reason: tool_calls → tool_use. + if final.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason = %q, want tool_use", final.StopReason) + } + // Usage: prompt→input, completion→output. + if final.Usage == nil || final.Usage.InputTokens != 11 || final.Usage.OutputTokens != 8 { + t.Errorf("usage = %+v, want input=11 output=8", final.Usage) + } + // Response identity. + if final.ResponseID != "chatcmpl-1" || final.ResponseModel != "gpt-4o" { + t.Errorf("response id/model = %q/%q", final.ResponseID, final.ResponseModel) + } + + // Content blocks: text first, then the tool call. + if len(final.Content) != 2 { + t.Fatalf("expected 2 content blocks, got %d: %+v", len(final.Content), final.Content) + } + txt, ok := final.Content[0].(agentcore.TextContent) + if !ok || txt.Text != "Let me check." { + t.Errorf("text block = %+v", final.Content[0]) + } + tool, ok := final.Content[1].(agentcore.ToolCallContent) + if !ok || tool.Name != "get_weather" || tool.ID != "call_1" { + t.Fatalf("tool block = %+v", final.Content[1]) + } + // tool_call arguments must have accumulated into valid JSON across fragments. + var args map[string]string + if err := json.Unmarshal(tool.Arguments, &args); err != nil { + t.Fatalf("tool arguments not valid JSON %q: %v", tool.Arguments, err) + } + if args["city"] != "SF" { + t.Errorf("tool arguments = %v, want city=SF", args) + } +} + +func TestOpenAIDecoderTextOnlyStop(t *testing.T) { + body := `data: {"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"Hello"}}]} + +data: {"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":" world"}}]} + +data: {"id":"c1","model":"gpt-4o","choices":[{"delta":{},"finish_reason":"stop"}]} + +data: {"choices":[],"usage":{"prompt_tokens":3,"completion_tokens":2}} + +data: [DONE] + +` + dec := NewOpenAIDecoder() + _, final := feedSSE(t, dec, body) + if final.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("stop reason = %q, want end_turn", final.StopReason) + } + if len(final.Content) != 1 { + t.Fatalf("expected 1 content block, got %d", len(final.Content)) + } + txt, ok := final.Content[0].(agentcore.TextContent) + if !ok || txt.Text != "Hello world" { + t.Errorf("text = %+v", final.Content[0]) + } +} + +func TestOpenAIDecoderLengthMapsToLength(t *testing.T) { + body := `data: {"id":"c","model":"m","choices":[{"delta":{"content":"truncated"}}]} + +data: {"id":"c","model":"m","choices":[{"delta":{},"finish_reason":"length"}]} + +data: [DONE] + +` + dec := NewOpenAIDecoder() + _, final := feedSSE(t, dec, body) + if final.StopReason != agentcore.StopReasonLength { + t.Errorf("length finish_reason must map to length, got %q", final.StopReason) + } +} + +// TestOpenAIDecoderInlineError verifies an inline error object becomes a decode +// error (which the transport turns into a terminal error event), never a panic. +func TestOpenAIDecoderInlineError(t *testing.T) { + dec := NewOpenAIDecoder() + _, err := dec.Decode([]byte(`{"error":{"type":"rate_limit_exceeded","message":"slow down"}}`)) + if err == nil { + t.Fatal("inline error object must return a decode error") + } + if !strings.Contains(err.Error(), "rate_limit_exceeded") { + t.Errorf("error should name the type, got %v", err) + } +} + +// TestOpenAIDecoderMalformedPayload verifies invalid JSON is a returned error +// (rides the stream as terminal error), not a panic. +func TestOpenAIDecoderMalformedPayload(t *testing.T) { + dec := NewOpenAIDecoder() + if _, err := dec.Decode([]byte(`{not json`)); err == nil { + t.Fatal("malformed payload must return an error") + } +} + +// TestOpenAIDecoderFinishFlushesPartial verifies a stream cut short (no +// finish_reason) still yields a done event on Finish, defaulting to end_turn. +func TestOpenAIDecoderFinishFlushesPartial(t *testing.T) { + body := `data: {"id":"c","model":"m","choices":[{"delta":{"content":"partial"}}]} + +` + dec := NewOpenAIDecoder() + events, final := feedSSE(t, dec, body) + if events[len(events)-1].EventKind() != StreamEventDone { + t.Fatalf("Finish must emit a terminal done event, got %v", eventKinds(events)) + } + if final.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("cut-short stream should default to end_turn, got %q", final.StopReason) + } + if len(final.Content) != 1 { + t.Fatalf("expected the partial text block, got %+v", final.Content) + } +} + +// TestOpenAIDecoderThroughTransport wires the decoder through the real transport +// pump against a recorded SSE server, exercising the full path including the +// [DONE] terminator handling. +func TestOpenAIDecoderThroughTransport(t *testing.T) { + srv := sseServer(t, openaiToolCallSSE) + defer srv.Close() + + stream, err := StreamRequest(context.Background(), TransportConfig{ + NewRequest: newReqFn(srv.URL), + Decoder: NewOpenAIDecoder(), + }) + if err != nil { + t.Fatalf("StreamRequest: %v", err) + } + var kinds []string + for ev := range stream.Events() { + kinds = append(kinds, ev.EventKind()) + } + final, resErr := stream.Result(context.Background()) + if resErr != nil { + t.Fatalf("result: %v", resErr) + } + if final.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason via transport = %q, want tool_use", final.StopReason) + } + if len(final.Content) != 2 { + t.Errorf("expected 2 content blocks via transport, got %d", len(final.Content)) + } + if kinds[len(kinds)-1] != StreamEventDone { + t.Errorf("stream must end with done, got %v", kinds) + } +} diff --git a/pigo/internal/provider/presets.go b/pigo/internal/provider/presets.go new file mode 100644 index 0000000..63f4738 --- /dev/null +++ b/pigo/internal/provider/presets.go @@ -0,0 +1,200 @@ +// This file defines the built-in preset catalog: a curated set of ready-to-use +// (provider, model) pairs a user can pick from without knowing each gateway's +// wire details (mirrors pi agent's preset provider/model picker). It covers the +// OpenAI-compatible gateways pigo ships — OpenRouter and NVIDIA NIM — plus a few +// local Ollama defaults. +// +// A preset binds a model id to the provider that serves it and that provider's +// default endpoint, so selecting a preset is enough to build a working Provider. +// The naive prefix-based mapping (ollama/…) still works for arbitrary ids; the +// preset catalog is the "menu" of vetted choices surfaced to the user. +// +// Security: presets carry no secrets. Each provider resolves its API key by name +// from the environment at request time (see auth.go); keys are never embedded +// here or logged. +package provider + +// PresetModel is one entry in the preset catalog: a model the user can select by +// id, the provider that serves it, and a short human label for the picker. +type PresetModel struct { + // Provider is the owning provider name (e.g. "openrouter", "nvidia"). + Provider string + // ID is the model id passed to the provider (e.g. "openai/gpt-4o"). + ID string + // DisplayName is a friendly label shown in the picker; falls back to ID. + DisplayName string +} + +// Label returns the display label for a preset, falling back to the id. +func (p PresetModel) Label() string { + if p.DisplayName != "" { + return p.DisplayName + } + return p.ID +} + +// PresetProviders lists the providers the preset catalog draws from, with the +// environment variable each expects its API key in (referenced by name only, +// never a value). Order is the display order in the picker. +var PresetProviders = []struct { + Name string + EnvVar string +}{ + {Name: "openrouter", EnvVar: "OPENROUTER_API_KEY"}, + {Name: "nvidia", EnvVar: "NVIDIA_API_KEY"}, + {Name: "deepseek", EnvVar: "DEEPSEEK_API_KEY"}, + {Name: "groq", EnvVar: "GROQ_API_KEY"}, + {Name: "xai", EnvVar: "XAI_API_KEY"}, + {Name: "cerebras", EnvVar: "CEREBRAS_API_KEY"}, + {Name: "mistral", EnvVar: "MISTRAL_API_KEY"}, + {Name: "moonshotai", EnvVar: "MOONSHOT_API_KEY"}, + {Name: "zai", EnvVar: "ZAI_API_KEY"}, + {Name: "fireworks", EnvVar: "FIREWORKS_API_KEY"}, + {Name: "together", EnvVar: "TOGETHER_API_KEY"}, + {Name: "minimax", EnvVar: "MINIMAX_API_KEY"}, + {Name: "xiaomi", EnvVar: "XIAOMI_API_KEY"}, + {Name: "qianfan", EnvVar: "QIANFAN_API_KEY"}, + {Name: "volcengine", EnvVar: "ARK_API_KEY"}, + {Name: "dashscope", EnvVar: "DASHSCOPE_API_KEY"}, + {Name: "hunyuan", EnvVar: "HUNYUAN_API_KEY"}, + {Name: "ollama", EnvVar: ""}, // local, no key +} + +// PresetCatalog is the built-in curated list of selectable models, grouped by +// provider in the order PresetProviders declares. These ids are the ones a user +// can `/model ` into or pick from `/models`; the list is representative, not +// exhaustive — any valid id for a known provider still works. +var PresetCatalog = []PresetModel{ + // --- OpenRouter (routes to many upstreams via one OpenAI-compatible API) --- + {Provider: "openrouter", ID: "openai/gpt-4o", DisplayName: "GPT-4o (OpenRouter)"}, + {Provider: "openrouter", ID: "openai/gpt-4o-mini", DisplayName: "GPT-4o mini (OpenRouter)"}, + {Provider: "openrouter", ID: "anthropic/claude-3.5-sonnet", DisplayName: "Claude 3.5 Sonnet (OpenRouter)"}, + {Provider: "openrouter", ID: "anthropic/claude-3.7-sonnet", DisplayName: "Claude 3.7 Sonnet (OpenRouter)"}, + {Provider: "openrouter", ID: "google/gemini-2.0-flash-001", DisplayName: "Gemini 2.0 Flash (OpenRouter)"}, + {Provider: "openrouter", ID: "google/gemini-2.5-pro", DisplayName: "Gemini 2.5 Pro (OpenRouter)"}, + {Provider: "openrouter", ID: "meta-llama/llama-3.3-70b-instruct", DisplayName: "Llama 3.3 70B (OpenRouter)"}, + {Provider: "openrouter", ID: "deepseek/deepseek-chat", DisplayName: "DeepSeek V3 (OpenRouter)"}, + {Provider: "openrouter", ID: "deepseek/deepseek-r1", DisplayName: "DeepSeek R1 (OpenRouter)"}, + {Provider: "openrouter", ID: "qwen/qwen-2.5-72b-instruct", DisplayName: "Qwen 2.5 72B (OpenRouter)"}, + {Provider: "openrouter", ID: "mistralai/mistral-large", DisplayName: "Mistral Large (OpenRouter)"}, + {Provider: "openrouter", ID: "x-ai/grok-2-1212", DisplayName: "Grok 2 (OpenRouter)"}, + + // --- OpenRouter free tier (":free" ids are rate-limited but cost nothing) --- + {Provider: "openrouter", ID: "deepseek/deepseek-r1:free", DisplayName: "DeepSeek R1 · free (OpenRouter)"}, + {Provider: "openrouter", ID: "deepseek/deepseek-chat-v3-0324:free", DisplayName: "DeepSeek V3 · free (OpenRouter)"}, + {Provider: "openrouter", ID: "meta-llama/llama-3.3-70b-instruct:free", DisplayName: "Llama 3.3 70B · free (OpenRouter)"}, + {Provider: "openrouter", ID: "google/gemini-2.0-flash-exp:free", DisplayName: "Gemini 2.0 Flash · free (OpenRouter)"}, + {Provider: "openrouter", ID: "qwen/qwen-2.5-72b-instruct:free", DisplayName: "Qwen 2.5 72B · free (OpenRouter)"}, + {Provider: "openrouter", ID: "qwen/qwq-32b:free", DisplayName: "QwQ 32B · free (OpenRouter)"}, + {Provider: "openrouter", ID: "mistralai/mistral-small-3.1-24b-instruct:free", DisplayName: "Mistral Small 3.1 24B · free (OpenRouter)"}, + {Provider: "openrouter", ID: "meta-llama/llama-4-maverick:free", DisplayName: "Llama 4 Maverick · free (OpenRouter)"}, + + // --- NVIDIA NIM (hosted, OpenAI-compatible) --- + {Provider: "nvidia", ID: "meta/llama-3.3-70b-instruct", DisplayName: "Llama 3.3 70B (NVIDIA)"}, + {Provider: "nvidia", ID: "meta/llama-3.1-405b-instruct", DisplayName: "Llama 3.1 405B (NVIDIA)"}, + {Provider: "nvidia", ID: "deepseek-ai/deepseek-r1", DisplayName: "DeepSeek R1 (NVIDIA)"}, + {Provider: "nvidia", ID: "qwen/qwen2.5-coder-32b-instruct", DisplayName: "Qwen 2.5 Coder 32B (NVIDIA)"}, + {Provider: "nvidia", ID: "nvidia/llama-3.1-nemotron-70b-instruct", DisplayName: "Nemotron 70B (NVIDIA)"}, + {Provider: "nvidia", ID: "mistralai/mixtral-8x22b-instruct-v0.1", DisplayName: "Mixtral 8x22B (NVIDIA)"}, + // NVIDIA's hosted NIM endpoint is free to call with a build.nvidia.com key. + {Provider: "nvidia", ID: "meta/llama-3.1-8b-instruct", DisplayName: "Llama 3.1 8B (NVIDIA)"}, + {Provider: "nvidia", ID: "meta/llama-3.1-70b-instruct", DisplayName: "Llama 3.1 70B (NVIDIA)"}, + {Provider: "nvidia", ID: "deepseek-ai/deepseek-v3", DisplayName: "DeepSeek V3 (NVIDIA)"}, + {Provider: "nvidia", ID: "qwen/qwen2.5-7b-instruct", DisplayName: "Qwen 2.5 7B (NVIDIA)"}, + {Provider: "nvidia", ID: "google/gemma-2-9b-it", DisplayName: "Gemma 2 9B (NVIDIA)"}, + {Provider: "nvidia", ID: "microsoft/phi-3.5-mini-instruct", DisplayName: "Phi-3.5 Mini (NVIDIA)"}, + + // --- DeepSeek (direct, OpenAI-compatible; ids from pi deepseek.models.ts) --- + {Provider: "deepseek", ID: "deepseek-v4-flash", DisplayName: "DeepSeek V4 Flash"}, + {Provider: "deepseek", ID: "deepseek-v4-pro", DisplayName: "DeepSeek V4 Pro"}, + + // --- Groq (fast inference; ids from pi groq.models.ts) --- + {Provider: "groq", ID: "llama-3.3-70b-versatile", DisplayName: "Llama 3.3 70B (Groq)"}, + {Provider: "groq", ID: "openai/gpt-oss-120b", DisplayName: "GPT OSS 120B (Groq)"}, + {Provider: "groq", ID: "qwen/qwen3-32b", DisplayName: "Qwen3 32B (Groq)"}, + + // --- xAI Grok (ids from pi xai.models.ts) --- + {Provider: "xai", ID: "grok-4.5", DisplayName: "Grok 4.5"}, + {Provider: "xai", ID: "grok-4.3", DisplayName: "Grok 4.3"}, + + // --- Cerebras (fast inference; ids from pi cerebras.models.ts) --- + {Provider: "cerebras", ID: "gpt-oss-120b", DisplayName: "GPT OSS 120B (Cerebras)"}, + {Provider: "cerebras", ID: "zai-glm-4.7", DisplayName: "Z.AI GLM-4.7 (Cerebras)"}, + {Provider: "cerebras", ID: "gemma-4-31b", DisplayName: "Gemma 4 31B (Cerebras)"}, + + // --- Mistral (ids from pi mistral.models.ts) --- + {Provider: "mistral", ID: "mistral-large-latest", DisplayName: "Mistral Large (latest)"}, + {Provider: "mistral", ID: "mistral-medium-latest", DisplayName: "Mistral Medium (latest)"}, + {Provider: "mistral", ID: "codestral-latest", DisplayName: "Codestral (latest)"}, + {Provider: "mistral", ID: "devstral-medium-latest", DisplayName: "Devstral Medium (latest)"}, + + // --- Moonshot AI Kimi (ids from pi moonshotai.models.ts) --- + {Provider: "moonshotai", ID: "kimi-k2-thinking", DisplayName: "Kimi K2 Thinking"}, + {Provider: "moonshotai", ID: "kimi-k2.6", DisplayName: "Kimi K2.6"}, + {Provider: "moonshotai", ID: "kimi-k3", DisplayName: "Kimi K3"}, + + // --- Z.AI GLM (ids from pi zai.models.ts) --- + {Provider: "zai", ID: "glm-4.7", DisplayName: "GLM-4.7"}, + {Provider: "zai", ID: "glm-5.1", DisplayName: "GLM-5.1"}, + {Provider: "zai", ID: "glm-5.2", DisplayName: "GLM-5.2"}, + + // --- Fireworks (ids from pi fireworks.models.ts) --- + {Provider: "fireworks", ID: "accounts/fireworks/models/deepseek-v4-pro", DisplayName: "DeepSeek V4 Pro (Fireworks)"}, + {Provider: "fireworks", ID: "accounts/fireworks/models/gpt-oss-120b", DisplayName: "GPT OSS 120B (Fireworks)"}, + {Provider: "fireworks", ID: "accounts/fireworks/models/kimi-k2p7-code", DisplayName: "Kimi K2.7 Code (Fireworks)"}, + + // --- Together AI (ids from pi together.models.ts) --- + {Provider: "together", ID: "deepseek-ai/DeepSeek-V4-Pro", DisplayName: "DeepSeek V4 Pro (Together)"}, + {Provider: "together", ID: "Qwen/Qwen3.7-Max", DisplayName: "Qwen3.7 Max (Together)"}, + {Provider: "together", ID: "meta-llama/Llama-3.3-70B-Instruct-Turbo", DisplayName: "Llama 3.3 70B Turbo (Together)"}, + + // --- MiniMax (Anthropic-protocol; ids from pi minimax.models.ts) --- + {Provider: "minimax", ID: "MiniMax-M2.7", DisplayName: "MiniMax-M2.7"}, + {Provider: "minimax", ID: "MiniMax-M3", DisplayName: "MiniMax-M3"}, + + // --- Xiaomi MiMo (ids from pi xiaomi.models.ts) --- + {Provider: "xiaomi", ID: "mimo-v2-pro", DisplayName: "MiMo-V2-Pro"}, + {Provider: "xiaomi", ID: "mimo-v2.5", DisplayName: "MiMo-V2.5"}, + {Provider: "xiaomi", ID: "mimo-v2.5-pro", DisplayName: "MiMo-V2.5-Pro"}, + + // --- Baidu AI Cloud Qianfan (OpenAI-compatible; ERNIE family) --- + {Provider: "qianfan", ID: "ernie-4.5-turbo-32k", DisplayName: "ERNIE 4.5 Turbo (Baidu Qianfan)"}, + + // --- ByteDance Volcengine Ark (OpenAI-compatible; Doubao family) --- + // Some Ark models require an "inference endpoint ID (endpoint id)" instead of a model + // name — use --base-url / -m to target those; this preset uses a model name. + {Provider: "volcengine", ID: "doubao-seed-1-6", DisplayName: "Doubao Seed 1.6 (Volcengine Ark)"}, + + // --- Alibaba Cloud DashScope (OpenAI-compatible; Qwen family) --- + {Provider: "dashscope", ID: "qwen-max", DisplayName: "Qwen Max (Alibaba DashScope)"}, + + // --- Tencent Hunyuan (OpenAI-compatible) --- + {Provider: "hunyuan", ID: "hunyuan-turbos-latest", DisplayName: "Hunyuan TurboS (Tencent Hunyuan)"}, + + // --- Ollama (local, no API key) --- + {Provider: "ollama", ID: "ollama/llama3.3", DisplayName: "Llama 3.3 (local Ollama)"}, + {Provider: "ollama", ID: "ollama/qwen2.5-coder", DisplayName: "Qwen 2.5 Coder (local Ollama)"}, +} + +// LookupPreset returns the preset entry for a model id, if the id is in the +// catalog. Used to resolve a selected id to its owning provider. +func LookupPreset(id string) (PresetModel, bool) { + for _, p := range PresetCatalog { + if p.ID == id { + return p, true + } + } + return PresetModel{}, false +} + +// PresetsByProvider returns the presets served by a given provider name, in +// catalog order. +func PresetsByProvider(providerName string) []PresetModel { + var out []PresetModel + for _, p := range PresetCatalog { + if p.Provider == providerName { + out = append(out, p) + } + } + return out +} diff --git a/pigo/internal/provider/presets_catalog_test.go b/pigo/internal/provider/presets_catalog_test.go new file mode 100644 index 0000000..2917e56 --- /dev/null +++ b/pigo/internal/provider/presets_catalog_test.go @@ -0,0 +1,118 @@ +package provider + +import "testing" + +// TestPresetProvidersIncludeNewProviders verifies the curated preset provider +// list gained the expanded set of gateways, each paired with the correct API-key +// environment variable (referenced by name only). +func TestPresetProvidersIncludeNewProviders(t *testing.T) { + byName := make(map[string]string, len(PresetProviders)) + for _, p := range PresetProviders { + byName[p.Name] = p.EnvVar + } + + want := map[string]string{ + "deepseek": "DEEPSEEK_API_KEY", + "groq": "GROQ_API_KEY", + "xai": "XAI_API_KEY", + "cerebras": "CEREBRAS_API_KEY", + "mistral": "MISTRAL_API_KEY", + "moonshotai": "MOONSHOT_API_KEY", + "zai": "ZAI_API_KEY", + "fireworks": "FIREWORKS_API_KEY", + "together": "TOGETHER_API_KEY", + "minimax": "MINIMAX_API_KEY", + "xiaomi": "XIAOMI_API_KEY", + } + for name, env := range want { + got, ok := byName[name] + if !ok { + t.Errorf("PresetProviders missing provider %q", name) + continue + } + if got != env { + t.Errorf("provider %q env var = %q, want %q", name, got, env) + } + } + + // Every preset provider must be a known provider in the central registry, so + // selecting a preset can always be resolved to a working Provider. + for _, p := range PresetProviders { + if p.Name == "ollama" { + continue // local pseudo-provider, not in the registry + } + if _, ok := LookupProviderSpec(p.Name); !ok { + t.Errorf("preset provider %q has no ProviderSpec in the registry", p.Name) + } + } +} + +// TestPresetCatalogCountsPerProvider asserts each expanded provider contributes +// the expected number of curated entries. +func TestPresetCatalogCountsPerProvider(t *testing.T) { + wantAtLeast := map[string]int{ + "deepseek": 2, + "groq": 3, + "xai": 2, + "cerebras": 3, + "mistral": 4, + "moonshotai": 3, + "zai": 3, + "fireworks": 3, + "together": 3, + "minimax": 2, + "xiaomi": 3, + } + for provider, min := range wantAtLeast { + got := PresetsByProvider(provider) + if len(got) < min { + t.Errorf("PresetsByProvider(%q) returned %d entries, want >= %d", provider, len(got), min) + } + for _, p := range got { + if p.Provider != provider { + t.Errorf("PresetsByProvider(%q) returned entry for %q", provider, p.Provider) + } + if p.ID == "" { + t.Errorf("PresetsByProvider(%q) returned entry with empty ID", provider) + } + } + } +} + +// TestLookupPresetNewEntries verifies representative new model ids resolve to the +// correct owning provider and carry a non-empty label. +func TestLookupPresetNewEntries(t *testing.T) { + cases := []struct { + id string + provider string + }{ + {"deepseek-v4-pro", "deepseek"}, + {"llama-3.3-70b-versatile", "groq"}, + {"grok-4.5", "xai"}, + {"zai-glm-4.7", "cerebras"}, + {"mistral-large-latest", "mistral"}, + {"kimi-k2-thinking", "moonshotai"}, + {"glm-5.1", "zai"}, + {"accounts/fireworks/models/gpt-oss-120b", "fireworks"}, + {"deepseek-ai/DeepSeek-V4-Pro", "together"}, + {"MiniMax-M3", "minimax"}, + {"mimo-v2.5-pro", "xiaomi"}, + } + for _, tc := range cases { + p, ok := LookupPreset(tc.id) + if !ok { + t.Errorf("LookupPreset(%q) not found", tc.id) + continue + } + if p.Provider != tc.provider { + t.Errorf("LookupPreset(%q).Provider = %q, want %q", tc.id, p.Provider, tc.provider) + } + if p.Label() == "" { + t.Errorf("LookupPreset(%q).Label() is empty", tc.id) + } + } + + if _, ok := LookupPreset("this-model-does-not-exist"); ok { + t.Error("LookupPreset returned ok for an unknown id") + } +} diff --git a/pigo/internal/provider/presets_test.go b/pigo/internal/provider/presets_test.go new file mode 100644 index 0000000..2b822ee --- /dev/null +++ b/pigo/internal/provider/presets_test.go @@ -0,0 +1,66 @@ +package provider + +// Tests for the preset provider/model catalog (mirrors pi agent's preset picker): +// LookupPreset resolves a catalog id to its owning provider, PresetsByProvider +// groups by provider, and every preset must name a provider that has a known +// credential env var (or be the local, keyless Ollama). + +import "testing" + +// TestLookupPresetResolvesProvider verifies a catalog id resolves to its +// declared provider, and an unknown id does not. +func TestLookupPresetResolvesProvider(t *testing.T) { + p, ok := LookupPreset("meta/llama-3.3-70b-instruct") + if !ok { + t.Fatal("expected NVIDIA llama preset to be in the catalog") + } + if p.Provider != "nvidia" { + t.Errorf("provider = %q, want nvidia", p.Provider) + } + if _, ok := LookupPreset("definitely/not-a-preset"); ok { + t.Error("unknown id must not resolve to a preset") + } +} + +// TestPresetsByProviderGroups verifies each provider surfaces at least one +// preset and that the returned entries all belong to that provider. +func TestPresetsByProviderGroups(t *testing.T) { + for _, name := range []string{"openrouter", "nvidia", "ollama"} { + got := PresetsByProvider(name) + if len(got) == 0 { + t.Errorf("provider %q has no presets", name) + } + for _, m := range got { + if m.Provider != name { + t.Errorf("PresetsByProvider(%q) returned entry for %q", name, m.Provider) + } + } + } +} + +// TestPresetProvidersHaveCredentialMapping verifies every non-local preset +// provider has API-key env vars in the provider registry (the single source of +// truth), so a selected preset can actually resolve a credential. Ollama is +// local and keyless. +func TestPresetProvidersHaveCredentialMapping(t *testing.T) { + for _, pv := range PresetProviders { + if pv.Name == "ollama" { + continue + } + spec, ok := LookupProviderSpec(pv.Name) + if !ok || len(spec.EnvVars) == 0 { + t.Errorf("preset provider %q has no credential env var mapping", pv.Name) + } + } +} + +// TestPresetLabelFallsBackToID verifies Label uses DisplayName when set and +// falls back to the id otherwise. +func TestPresetLabelFallsBackToID(t *testing.T) { + if got := (PresetModel{ID: "x", DisplayName: "X"}).Label(); got != "X" { + t.Errorf("Label = %q, want X", got) + } + if got := (PresetModel{ID: "x"}).Label(); got != "x" { + t.Errorf("Label = %q, want x", got) + } +} diff --git a/pigo/internal/provider/protocol.go b/pigo/internal/provider/protocol.go new file mode 100644 index 0000000..26c909c --- /dev/null +++ b/pigo/internal/provider/protocol.go @@ -0,0 +1,77 @@ +package provider + +// Protocol normalization (US-001, #538). The user-facing --protocol / protocol +// value accepts three OpenAI wire variants in addition to anthropic: +// +// openai → Chat Completions (POST {base_url}/chat/completions) +// openai/chat → Chat Completions (alias of "openai") +// openai/resp_api → Responses API (POST {base_url}/responses) +// anthropic → Anthropic Messages +// "" → unset; downstream falls back to model-id heuristics +// +// NormalizeProtocol collapses these into a small set of canonical internal +// selectors so ResolveProvider (#543) can switch on chat vs resp_api without +// re-parsing surface syntax. "openai" and "openai/chat" both normalize to +// ProtocolOpenAI, keeping the existing Chat Completions path byte-for-byte +// unchanged; only "openai/resp_api" produces the new selector. + +import ( + "fmt" + "strings" +) + +// ProtocolOpenAIResponses is the canonical selector for the OpenAI Responses +// API wire format (POST {base_url}/responses). It is distinct from +// ProtocolOpenAI (Chat Completions) so ResolveProvider can route to the +// SDK-based Responses driver. +const ProtocolOpenAIResponses = "openai/resp_api" + +// NormalizeProtocol maps a raw --protocol / protocol value to a canonical +// internal selector. Input is trimmed and lower-cased before matching. An empty +// value stays empty (unset → model-id heuristics). Recognized values normalize +// to ProtocolOpenAI, ProtocolOpenAIResponses, or ProtocolAnthropic. Any other +// value is an error naming the accepted set, so a typo surfaces to the caller +// for exit-code mapping instead of silently falling through. +func NormalizeProtocol(raw string) (string, error) { + switch strings.ToLower(strings.TrimSpace(raw)) { + case "": + return "", nil + case ProtocolOpenAI, "openai/chat": + return ProtocolOpenAI, nil + case ProtocolOpenAIResponses: + return ProtocolOpenAIResponses, nil + case ProtocolAnthropic: + return ProtocolAnthropic, nil + default: + return "", fmt.Errorf("unknown --protocol %q (want openai|openai/chat|openai/resp_api|anthropic)", raw) + } +} + +// ProtocolLabel maps a raw --protocol value to the human-facing label shown in +// the startup banner's Protocol row, so the displayed wire format matches what +// pigo actually speaks. It differs from NormalizeProtocol in one deliberate way: +// the bare "openai" input is surfaced as "openai/chat", making the Chat +// Completions variant explicit rather than ambiguous. "openai/resp_api" and +// "anthropic" pass through as themselves. +// +// An empty input returns empty (the banner then falls back to "—" or the +// provider name, so an unset protocol on a named/inferred provider is not +// mislabeled). An unrecognized value returns the trimmed input verbatim — the +// label is presentation-only and must never fail; a real typo is already +// rejected upstream by NormalizeProtocol during resolution. +func ProtocolLabel(raw string) string { + canonical, err := NormalizeProtocol(raw) + if err != nil { + return strings.TrimSpace(raw) + } + switch canonical { + case ProtocolOpenAI: + return "openai/chat" + case ProtocolOpenAIResponses: + return ProtocolOpenAIResponses + case ProtocolAnthropic: + return ProtocolAnthropic + default: + return "" + } +} diff --git a/pigo/internal/provider/protocol_test.go b/pigo/internal/provider/protocol_test.go new file mode 100644 index 0000000..f846589 --- /dev/null +++ b/pigo/internal/provider/protocol_test.go @@ -0,0 +1,55 @@ +package provider + +import ( + "strings" + "testing" +) + +func TestNormalizeProtocol(t *testing.T) { + tests := []struct { + name string + in string + want string + wantErr bool + }{ + {"empty stays empty", "", "", false}, + {"openai", "openai", ProtocolOpenAI, false}, + {"openai/chat aliases openai", "openai/chat", ProtocolOpenAI, false}, + {"openai/resp_api distinct", "openai/resp_api", ProtocolOpenAIResponses, false}, + {"anthropic unchanged", "anthropic", ProtocolAnthropic, false}, + {"case-insensitive", "OpenAI/Resp_API", ProtocolOpenAIResponses, false}, + {"trimmed", " openai/chat ", ProtocolOpenAI, false}, + {"unknown rejected", "openai/foo", "", true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got, err := NormalizeProtocol(tc.in) + if tc.wantErr { + if err == nil { + t.Fatalf("NormalizeProtocol(%q) = %q, want error", tc.in, got) + } + return + } + if err != nil { + t.Fatalf("NormalizeProtocol(%q) unexpected error: %v", tc.in, err) + } + if got != tc.want { + t.Errorf("NormalizeProtocol(%q) = %q, want %q", tc.in, got, tc.want) + } + }) + } +} + +// The rejection message must name every accepted value so a user with a typo +// can self-correct without reading source. +func TestNormalizeProtocolErrorNamesAcceptedValues(t *testing.T) { + _, err := NormalizeProtocol("openai/foo") + if err == nil { + t.Fatal("expected error for unknown protocol") + } + for _, want := range []string{"openai", "openai/chat", "openai/resp_api", "anthropic"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error %q missing accepted value %q", err.Error(), want) + } + } +} diff --git a/pigo/internal/provider/provider.go b/pigo/internal/provider/provider.go new file mode 100644 index 0000000..10164c4 --- /dev/null +++ b/pigo/internal/provider/provider.go @@ -0,0 +1,118 @@ +// This file defines the provider streaming abstraction (US-003/US-007 base): +// the StreamFn contract, the per-delta AssistantMessageEvent set, and the +// AssistantMessageEventStream (a specialization of EventStream) that a provider +// pushes deltas onto while yielding a final AssistantMessage. +// +// Contract (FR-13): a StreamFn never expresses a request failure by returning +// an error. Runtime failures are encoded as an error event plus a terminal +// assistant message (stopReason=error/aborted + errorMessage). The returned +// error is reserved for the earliest "could not even build the stream" case. +package provider + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// AssistantMessageEvent is the sealed interface for provider stream deltas. The +// loop dispatches on EventKind; the raw event is also surfaced to consumers via +// MessageUpdateEvent.AssistantMessageEvent. +type AssistantMessageEvent interface { + isAssistantMessageEvent() + // EventKind returns the delta discriminant. + EventKind() string +} + +// AssistantMessageEvent kinds. +const ( + StreamEventStart = "start" + StreamEventText = "text" + StreamEventThinking = "thinking" + StreamEventToolCall = "toolcall" + StreamEventDone = "done" + StreamEventError = "error" +) + +// StreamStartEvent carries the initial (usually empty) partial message. +type StreamStartEvent struct{ Partial agentcore.AssistantMessage } + +// StreamTextEvent carries the partial message after a text delta. +type StreamTextEvent struct{ Partial agentcore.AssistantMessage } + +// StreamThinkingEvent carries the partial after a thinking delta. +type StreamThinkingEvent struct{ Partial agentcore.AssistantMessage } + +// StreamToolCallEvent carries the partial after a tool-call delta. +type StreamToolCallEvent struct{ Partial agentcore.AssistantMessage } + +// StreamDoneEvent is the terminal success event; Message is the final response. +type StreamDoneEvent struct{ Message agentcore.AssistantMessage } + +// StreamErrorEvent is the terminal failure event; Message carries the terminal +// assistant message (stopReason=error/aborted + errorMessage). +type StreamErrorEvent struct { + Message agentcore.AssistantMessage + Err error +} + +func (StreamStartEvent) isAssistantMessageEvent() {} +func (StreamTextEvent) isAssistantMessageEvent() {} +func (StreamThinkingEvent) isAssistantMessageEvent() {} +func (StreamToolCallEvent) isAssistantMessageEvent() {} +func (StreamDoneEvent) isAssistantMessageEvent() {} +func (StreamErrorEvent) isAssistantMessageEvent() {} + +func (StreamStartEvent) EventKind() string { return StreamEventStart } +func (StreamTextEvent) EventKind() string { return StreamEventText } +func (StreamThinkingEvent) EventKind() string { return StreamEventThinking } +func (StreamToolCallEvent) EventKind() string { return StreamEventToolCall } +func (StreamDoneEvent) EventKind() string { return StreamEventDone } +func (StreamErrorEvent) EventKind() string { return StreamEventError } + +// AssistantMessageEventStream is the provider-level stream: deltas of type +// AssistantMessageEvent with a final AssistantMessage result. isComplete fires +// on done/error; extractResult takes the terminal event's message. +type AssistantMessageEventStream = agentcore.EventStream[AssistantMessageEvent, agentcore.AssistantMessage] + +// NewAssistantMessageEventStream builds a provider stream wired with the +// done/error completion callbacks. +func NewAssistantMessageEventStream(buffer int) *AssistantMessageEventStream { + s := agentcore.NewEventStream[AssistantMessageEvent, agentcore.AssistantMessage](buffer) + s.IsComplete = func(e AssistantMessageEvent) bool { + k := e.EventKind() + return k == StreamEventDone || k == StreamEventError + } + s.ExtractResult = func(e AssistantMessageEvent) agentcore.AssistantMessage { + switch ev := e.(type) { + case StreamDoneEvent: + return ev.Message + case StreamErrorEvent: + return ev.Message + default: + return agentcore.AssistantMessage{} + } + } + return s +} + +// LlmContext is the shaped request handed to a StreamFn: the system prompt, the +// LLM-bound messages (UI-only messages already filtered), and the tools. +type LlmContext struct { + SystemPrompt string + Messages agentcore.MessageList + Tools []agentcore.AgentTool +} + +// StreamConfig carries per-request settings for a StreamFn. +type StreamConfig struct { + APIKey string + ThinkingLevel agentcore.ThinkingLevel + // Extra holds provider-specific options; opaque to the loop. + Extra map[string]any +} + +// StreamFn produces a provider stream for a model + shaped context. Per the +// contract it returns an error only for early "cannot build the stream" +// failures; all runtime failures ride the returned stream as error events. +type StreamFn func(ctx context.Context, model string, llm LlmContext, cfg StreamConfig) (*AssistantMessageEventStream, error) diff --git a/pigo/internal/provider/provider_interface.go b/pigo/internal/provider/provider_interface.go new file mode 100644 index 0000000..11944a8 --- /dev/null +++ b/pigo/internal/provider/provider_interface.go @@ -0,0 +1,79 @@ +// This file defines the unified Provider interface and its dual failure model +// (US-007). A Provider turns a CompletionRequest into a stream of +// AssistantMessageEvents. Failures follow the same contract as StreamFn (FR-13): +// +// - "cannot even build the stream" (bad config, missing model) → returned error. +// - any runtime failure once streaming has begun → a terminal StreamErrorEvent +// carrying an assistant message with stopReason=error/aborted, after which +// the stream is closed. It is never a returned error. +// +// Model carries provider-agnostic capability metadata so the loop and UI can +// reason about a model without knowing the concrete provider. +package provider + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// Model is provider-agnostic metadata describing a single model's identity and +// capabilities. Providers construct these; the loop/UI consume them. +type Model struct { + // Provider is the provider name (e.g. "anthropic", "openai"). + Provider string `json:"provider"` + // ID is the provider-specific model id (e.g. "claude-opus-4-8"). + ID string `json:"id"` + // DisplayName is a human-friendly label; falls back to ID when empty. + DisplayName string `json:"displayName,omitempty"` + // ContextWindow is the maximum input+output token window, 0 if unknown. + ContextWindow int `json:"contextWindow,omitempty"` + // MaxOutputTokens is the max tokens the model may emit per response, 0 if + // unknown. + MaxOutputTokens int `json:"maxOutputTokens,omitempty"` + // SupportsThinking reports whether the model exposes a reasoning/thinking + // channel. + SupportsThinking bool `json:"supportsThinking,omitempty"` + // SupportsTools reports whether the model can call tools. + SupportsTools bool `json:"supportsTools,omitempty"` + // SupportsImages reports whether the model accepts image (multimodal) input. + // When false, an image block in the request is reported as a hard error + // rather than silently dropped, so the user learns the model cannot see it. + SupportsImages bool `json:"supportsImages,omitempty"` + // ThinkingLevels maps unified thinking levels to this model's wire values. + // nil when the model does not support thinking (decision #10). + ThinkingLevels agentcore.ThinkingLevelMap `json:"-"` +} + +// CompletionRequest is the provider-agnostic input to StreamCompletion: the +// model id, the shaped LLM context, and per-request options. +type CompletionRequest struct { + // Model is the provider-specific model id to complete against. + Model string + // Context is the shaped request (system prompt, LLM-bound messages, tools). + Context LlmContext + // Config carries per-request options (API key, thinking level, extras). + Config StreamConfig +} + +// Provider is the unified streaming interface implemented by every backend. It +// hides per-vendor differences behind a single AssistantMessageEvent stream. +type Provider interface { + // Name returns the provider's identifier (matches Model.Provider). + Name() string + // Models lists the models this provider can serve. + Models() []Model + // StreamCompletion streams a completion for req. Per the dual failure model + // it returns an error only for the earliest "cannot build the stream" case; + // all runtime failures ride the returned stream as a terminal error event. + StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) +} + +// StreamFnFromProvider adapts a Provider to the loop's StreamFn contract so a +// Provider can drive streamAssistantResponse directly. The two failure models +// are identical, so the adaptation is a straight delegation. +func StreamFnFromProvider(p Provider) StreamFn { + return func(ctx context.Context, model string, llm LlmContext, cfg StreamConfig) (*AssistantMessageEventStream, error) { + return p.StreamCompletion(ctx, CompletionRequest{Model: model, Context: llm, Config: cfg}) + } +} diff --git a/pigo/internal/provider/provider_interface_test.go b/pigo/internal/provider/provider_interface_test.go new file mode 100644 index 0000000..33b408e --- /dev/null +++ b/pigo/internal/provider/provider_interface_test.go @@ -0,0 +1,87 @@ +package provider + +import ( + "context" + "errors" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// fakeProvider is a minimal Provider for interface tests. +type fakeProvider struct { + name string + models []Model + buildErr error + events []AssistantMessageEvent +} + +func (p fakeProvider) Name() string { return p.name } +func (p fakeProvider) Models() []Model { return p.models } +func (p fakeProvider) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) { + if p.buildErr != nil { + return nil, p.buildErr + } + s := NewAssistantMessageEventStream(0) + go func() { + for _, ev := range p.events { + if err := s.Emit(ctx, ev); err != nil { + s.SetError(err) + break + } + } + s.Close() + }() + return s, nil +} + +func TestProviderEarlyBuildFailureReturnsError(t *testing.T) { + p := fakeProvider{name: "test", buildErr: errors.New("no such model")} + _, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "ghost"}) + if err == nil { + t.Fatal("early build failure must return an error") + } +} + +func TestProviderRuntimeFailureRidesStream(t *testing.T) { + errMsg := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonError, ErrorMessage: "upstream 500"} + p := fakeProvider{ + name: "test", + events: []AssistantMessageEvent{StreamErrorEvent{Message: errMsg}}, + } + stream, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"}) + if err != nil { + t.Fatalf("runtime failure must NOT be a returned error: %v", err) + } + final, resErr := stream.Result(context.Background()) + if resErr != nil { + t.Fatalf("stream result error: %v", resErr) + } + if final.StopReason != agentcore.StopReasonError || final.ErrorMessage != "upstream 500" { + t.Errorf("terminal error message wrong: %+v", final) + } +} + +func TestStreamFnFromProviderDelegates(t *testing.T) { + done := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn} + p := fakeProvider{name: "test", events: []AssistantMessageEvent{StreamDoneEvent{Message: done}}} + fn := StreamFnFromProvider(p) + stream, err := fn(context.Background(), "m", LlmContext{}, StreamConfig{}) + if err != nil { + t.Fatalf("delegation error: %v", err) + } + final, _ := stream.Result(context.Background()) + if final.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("delegated stream result wrong: %+v", final) + } +} + +func TestModelMetadata(t *testing.T) { + m := Model{Provider: "anthropic", ID: "claude-opus-4-8", SupportsThinking: true, ContextWindow: 200000} + if m.Provider != "anthropic" || m.ID != "claude-opus-4-8" { + t.Errorf("model identity wrong: %+v", m) + } + if !m.SupportsThinking || m.ContextWindow != 200000 { + t.Errorf("model capability wrong: %+v", m) + } +} diff --git a/pigo/internal/provider/providers.go b/pigo/internal/provider/providers.go new file mode 100644 index 0000000..0257790 --- /dev/null +++ b/pigo/internal/provider/providers.go @@ -0,0 +1,775 @@ +// This file implements the first concrete Providers (US-013): Bedrock, +// OpenRouter, and Ollama. Each satisfies the Provider interface by building an +// *http.Request and delegating to the shared transport (StreamRequest) with a +// provider-appropriate Decoder — no bespoke HTTP/SSE handling per provider. +// +// Two backing driver shapes cover all three: +// +// - openAICompatDriver — POSTs to {baseURL}/chat/completions in the OpenAI +// Chat Completions wire format, decoded by OpenAIDecoder. OpenRouter and +// Ollama are instances of it; it is the generic OpenAI-compatible layer, +// reusable for any gateway (Groq, together, local servers, …). +// - anthropicCompatDriver — POSTs the Anthropic Messages wire format, decoded +// by AnthropicDecoder. Bedrock rides this: Anthropic-on-Bedrock speaks the +// Messages API, so the decoder is reused wholesale. +// +// Failures follow the dual failure model (FR-13): only the earliest "cannot +// build the stream" case (missing key, bad request construction) is a returned +// error; every runtime failure rides the stream as a terminal error event, +// which StreamRequest already guarantees. +// +// Security (US-012 / US-026): API keys are referenced by provider name in any +// error; secret values are never logged or embedded in error text. +package provider + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// providerBaseURLs holds the default endpoint per built-in provider. A base URL +// is a transport concern (the decoders are URL-agnostic), so overriding it — +// e.g. pointing Ollama at a remote host — needs no decoder change. +const ( + openRouterBaseURL = "https://openrouter.ai/api/v1" + ollamaBaseURL = "http://localhost:11434/v1" + // nvidiaBaseURL is NVIDIA's hosted NIM endpoint. It speaks the OpenAI + // Chat Completions wire format, so it rides openAICompatDriver unchanged. + nvidiaBaseURL = "https://integrate.api.nvidia.com/v1" + // bedrockBaseURL is a placeholder default; real Bedrock endpoints are + // region-specific (bedrock-runtime..amazonaws.com) and supplied at + // construction. It is exported as a field so callers set the resolved URL. + bedrockBaseURL = "https://bedrock-runtime.us-east-1.amazonaws.com" + // anthropicBaseURL is the public Anthropic Messages API endpoint, the default + // for a --protocol=anthropic provider when no --base-url is given. + anthropicBaseURL = "https://api.anthropic.com/v1" + // anthropicAPIVersion is the required anthropic-version header value sent with + // every direct-Anthropic request. + anthropicAPIVersion = "2023-06-01" +) + +// --------------------------------------------------------------------------- +// OpenAI-compatible driver (OpenRouter, Ollama, and any OpenAI-compatible API). +// --------------------------------------------------------------------------- + +// openAICompatDriver is the shared backing for every OpenAI-compatible +// provider. It holds the provider identity, endpoint, model catalog, and the +// auth scheme; StreamCompletion builds the chat-completions request and hands +// it to the transport with a fresh OpenAIDecoder. +type openAICompatDriver struct { + name string + baseURL string + models []Model + // requiresAuth reports whether an Authorization: Bearer header is sent. + // Ollama (local) needs none; OpenRouter does. + requiresAuth bool + // extraHeaders are attached to every request (e.g. OpenRouter attribution). + extraHeaders map[string]string +} + +func (d *openAICompatDriver) Name() string { return d.name } +func (d *openAICompatDriver) Models() []Model { return d.models } + +// StreamCompletion builds the OpenAI Chat Completions request and streams it. +func (d *openAICompatDriver) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) { + if d.requiresAuth && strings.TrimSpace(req.Config.APIKey) == "" { + // Early "cannot build the stream": reference the provider, never a value. + return nil, fmt.Errorf("%s: missing API key", d.name) + } + if err := checkImageSupport(d.name, req.Model, d.models, req.Context.Messages); err != nil { + return nil, err + } + body, err := encodeOpenAIRequest(req) + if err != nil { + return nil, fmt.Errorf("%s: build request body: %w", d.name, err) + } + newReq := func(ctx context.Context) (*http.Request, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, d.baseURL+"/chat/completions", bytes.NewReader(body)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + if d.requiresAuth { + httpReq.Header.Set("Authorization", "Bearer "+req.Config.APIKey) + } + for k, v := range d.extraHeaders { + httpReq.Header.Set(k, v) + } + return httpReq, nil + } + return StreamRequest(ctx, TransportConfig{NewRequest: newReq, Decoder: NewOpenAIDecoder()}) +} + +// encodeOpenAIRequest serializes a CompletionRequest into an OpenAI Chat +// Completions JSON body with streaming enabled and usage requested. +func encodeOpenAIRequest(req CompletionRequest) ([]byte, error) { + msgs := make([]map[string]any, 0, len(req.Context.Messages)+1) + if sp := req.Context.SystemPrompt; sp != "" { + msgs = append(msgs, map[string]any{"role": "system", "content": sp}) + } + for _, m := range req.Context.Messages { + msgs = append(msgs, encodeOpenAIMessage(m)...) + } + body := map[string]any{ + "model": req.Model, + "messages": msgs, + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + } + // Reasoning effort: when a thinking level is requested, forward it as the + // OpenAI `reasoning_effort` field. Reasoning models (o-series, DeepSeek-R1, + // GLM-thinking, …) read this to open their reasoning channel; omitting it + // leaves them at their default and effectively disables extended reasoning. + if effort := openAIReasoningEffort(req.Config.ThinkingLevel); effort != "" { + body["reasoning_effort"] = effort + } + if tools := encodeOpenAITools(req.Context.Tools); len(tools) > 0 { + body["tools"] = tools + } + return json.Marshal(body) +} + +// openAIReasoningEffort maps the unified ThinkingLevel onto the OpenAI +// `reasoning_effort` wire value. "off"/"" yields "" (field omitted, default +// behavior preserved). OpenAI accepts minimal|low|medium|high; xhigh maps to +// high (the strongest supported value). +func openAIReasoningEffort(level agentcore.ThinkingLevel) string { + switch level { + case agentcore.ThinkingMinimal: + return "minimal" + case agentcore.ThinkingLow: + return "low" + case agentcore.ThinkingMedium: + return "medium" + case agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax: + return "high" + default: // off or unset + return "" + } +} + +// encodeOpenAIMessage maps one pigo message onto the OpenAI wire shape. An +// assistant message may expand to content + tool_calls in a single entry; a +// tool result becomes a role:"tool" entry keyed by tool_call_id. +func encodeOpenAIMessage(m agentcore.Message) []map[string]any { + switch msg := m.(type) { + case agentcore.UserMessage: + return []map[string]any{{"role": "user", "content": openAIUserContent(msg.Content)}} + case agentcore.CompactionMessage: + // A compaction checkpoint stands in for compacted history as user text. + u := msg.AsUserMessage() + return []map[string]any{{"role": "user", "content": openAIUserContent(u.Content)}} + case agentcore.AssistantMessage: + entry := map[string]any{"role": "assistant"} + text := agentcore.ContentToText(msg.Content) + var toolCalls []map[string]any + for _, c := range msg.Content { + if tc, ok := c.(agentcore.ToolCallContent); ok { + toolCalls = append(toolCalls, map[string]any{ + "id": tc.ID, + "type": "function", + "function": map[string]any{ + "name": tc.Name, + "arguments": string(tc.Arguments), + }, + }) + } + } + if len(toolCalls) > 0 { + entry["tool_calls"] = toolCalls + // With tool calls present, send content as JSON null when there is no + // accompanying text: an empty string trips strict gateways (e.g. vLLM) + // that expect null | non-empty for an assistant tool-call turn. + if text == "" { + entry["content"] = nil + } else { + entry["content"] = text + } + } else { + entry["content"] = text + } + return []map[string]any{entry} + case agentcore.ToolResultMessage: + return []map[string]any{{ + "role": "tool", + "tool_call_id": msg.ToolCallID, + "content": agentcore.ContentToText(msg.Content), + }} + default: + return nil + } +} + +// openAIUserContent shapes a user content list for the OpenAI wire. When there +// are no images it collapses to a plain string (the common case, and what most +// OpenAI-compatible gateways expect). When images are present it emits the +// multimodal array form: text parts plus image_url parts carrying a base64 data +// URI (data:;base64,). +func openAIUserContent(content agentcore.ContentList) any { + hasImage := false + for _, c := range content { + if _, ok := c.(agentcore.ImageContent); ok { + hasImage = true + break + } + } + if !hasImage { + return agentcore.ContentToText(content) + } + parts := make([]map[string]any, 0, len(content)) + for _, c := range content { + switch b := c.(type) { + case agentcore.TextContent: + if b.Text == "" { + continue + } + parts = append(parts, map[string]any{"type": "text", "text": b.Text}) + case agentcore.ImageContent: + parts = append(parts, map[string]any{ + "type": "image_url", + "image_url": map[string]any{ + "url": fmt.Sprintf("data:%s;base64,%s", b.MimeType, b.Data), + }, + }) + } + } + return parts +} + +// encodeOpenAITools maps AgentTools onto the OpenAI function-tool schema. +func encodeOpenAITools(tools []agentcore.AgentTool) []map[string]any { + if len(tools) == 0 { + return nil + } + out := make([]map[string]any, 0, len(tools)) + for _, t := range tools { + params := json.RawMessage(t.Schema()) + if len(params) == 0 { + params = json.RawMessage("{}") + } + out = append(out, map[string]any{ + "type": "function", + "function": map[string]any{ + "name": t.Name(), + "description": t.Description(), + "parameters": params, + }, + }) + } + return out +} + +// --------------------------------------------------------------------------- +// Anthropic-compatible driver (Bedrock). +// --------------------------------------------------------------------------- + +// anthropicCompatDriver backs Anthropic-wire providers that are not the direct +// Anthropic API — Bedrock being the case here. It POSTs the Messages wire +// format and decodes with AnthropicDecoder. +type anthropicCompatDriver struct { + name string + baseURL string + models []Model + // path is the endpoint path appended to baseURL (Bedrock's invoke path + // embeds the model id, so it is derived per request). + pathFor func(model string) string + // authHeader sets provider auth on the request (never logs the value). + authHeader func(req *http.Request, apiKey string) +} + +func (d *anthropicCompatDriver) Name() string { return d.name } +func (d *anthropicCompatDriver) Models() []Model { return d.models } + +// StreamCompletion builds the Anthropic Messages request and streams it, +// decoding with AnthropicDecoder. +func (d *anthropicCompatDriver) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) { + if strings.TrimSpace(req.Config.APIKey) == "" { + return nil, fmt.Errorf("%s: missing API key", d.name) + } + if err := checkImageSupport(d.name, req.Model, d.models, req.Context.Messages); err != nil { + return nil, err + } + body, err := encodeAnthropicRequest(req, d.models) + if err != nil { + return nil, fmt.Errorf("%s: build request body: %w", d.name, err) + } + path := "/messages" + if d.pathFor != nil { + path = d.pathFor(req.Model) + } + newReq := func(ctx context.Context) (*http.Request, error) { + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, d.baseURL+path, bytes.NewReader(body)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + if d.authHeader != nil { + d.authHeader(httpReq, req.Config.APIKey) + } + return httpReq, nil + } + return StreamRequest(ctx, TransportConfig{NewRequest: newReq, Decoder: NewAnthropicDecoder()}) +} + +// encodeAnthropicRequest serializes a CompletionRequest into an Anthropic +// Messages JSON body with streaming enabled. The system prompt is a top-level +// field; tool results and tool calls follow the Messages content-block shape. +func encodeAnthropicRequest(req CompletionRequest, models []Model) ([]byte, error) { + msgs := make([]map[string]any, 0, len(req.Context.Messages)) + all := req.Context.Messages + for i := 0; i < len(all); i++ { + m := all[i] + // Anthropic requires every tool_use in the preceding assistant message to + // be answered by tool_result blocks in a single immediately-following user + // message — splitting them across consecutive user turns yields the API + // error "tool_use ids were found without tool_result blocks immediately + // after". Coalesce runs of consecutive tool results into one turn. + if tm, ok := m.(agentcore.ToolResultMessage); ok { + blocks := []map[string]any{{ + "type": "tool_result", + "tool_use_id": tm.ToolCallID, + "content": agentcore.ContentToText(tm.Content), + }} + for i+1 < len(all) { + next, ok := all[i+1].(agentcore.ToolResultMessage) + if !ok { + break + } + i++ + blocks = append(blocks, map[string]any{ + "type": "tool_result", + "tool_use_id": next.ToolCallID, + "content": agentcore.ContentToText(next.Content), + }) + } + msgs = append(msgs, map[string]any{"role": "user", "content": blocks}) + continue + } + if enc := encodeAnthropicMessage(m); enc != nil { + msgs = append(msgs, enc) + } + } + body := map[string]any{ + "model": req.Model, + "messages": msgs, + "stream": true, + } + if sp := req.Context.SystemPrompt; sp != "" { + body["system"] = sp + } + maxTok := maxOutputTokensFor(req) + if maxTok <= 0 { + // Anthropic requires max_tokens. Prefer the model's declared cap; fall + // back to a coding-friendly default (4096 was too low and caused + // truncation/retry loops on longer edits). + maxTok = anthropicDefaultMaxTokens(req.Model, models) + } + // Extended thinking: when a thinking level is requested, enable the Anthropic + // thinking block with a budget derived from the level. Omitted for off/unset + // so non-thinking requests keep their prior shape. + if budget := anthropicThinkingBudget(req.Config.ThinkingLevel); budget > 0 { + // Anthropic counts thinking tokens toward max_tokens and requires + // budget_tokens < max_tokens (else a 400). Guarantee headroom for the + // visible reply by lifting max_tokens above the budget when the caller's + // cap is too low to fit both the reasoning and a real answer. + if minTok := budget + anthropicResponseHeadroom; maxTok < minTok { + maxTok = minTok + } + body["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget} + } + body["max_tokens"] = maxTok + if tools := encodeAnthropicTools(req.Context.Tools); len(tools) > 0 { + body["tools"] = tools + } + return json.Marshal(body) +} + +// anthropicResponseHeadroom is the token margin reserved for the visible reply +// on top of the thinking budget, so max_tokens always exceeds budget_tokens (an +// Anthropic hard requirement) with room left for a real answer. +const anthropicResponseHeadroom = 4096 + +// anthropicDefaultMaxTokens picks the max_tokens fallback when no explicit hint +// is given: the model's declared MaxOutputTokens if present in the driver's +// model catalog, otherwise 8192 (a coding-friendly default that avoids +// premature truncation while staying within common model caps). +func anthropicDefaultMaxTokens(model string, models []Model) int { + for _, m := range models { + if m.ID == model && m.MaxOutputTokens > 0 { + return m.MaxOutputTokens + } + } + return 8192 +} + +// anthropicThinkingBudget maps a unified ThinkingLevel onto an Anthropic +// thinking budget_tokens value. off/"" yields 0 (thinking block omitted). +func anthropicThinkingBudget(level agentcore.ThinkingLevel) int { + switch level { + case agentcore.ThinkingMinimal: + return 1024 + case agentcore.ThinkingLow: + return 2048 + case agentcore.ThinkingMedium: + return 8192 + case agentcore.ThinkingHigh: + return 16384 + case agentcore.ThinkingXHigh: + return 32768 + case agentcore.ThinkingMax: + return 65536 + default: // off or unset + return 0 + } +} + +// encodeAnthropicMessage maps one pigo message onto the Anthropic Messages +// wire shape. Assistant tool calls become tool_use blocks; tool results become +// a user message carrying a tool_result block (Anthropic's convention). +func encodeAnthropicMessage(m agentcore.Message) map[string]any { + switch msg := m.(type) { + case agentcore.UserMessage: + return map[string]any{"role": "user", "content": anthropicUserContent(msg.Content)} + case agentcore.CompactionMessage: + u := msg.AsUserMessage() + return map[string]any{"role": "user", "content": anthropicUserContent(u.Content)} + case agentcore.AssistantMessage: + var blocks []map[string]any + // Thinking blocks must precede tool_use in the same assistant turn: + // Anthropic extended-thinking requires the prior thinking block (and its + // signature) to be echoed back verbatim on tool-use turns, or the API + // rejects/degrades the request. Emit them first. + for _, c := range msg.Content { + if t, ok := c.(agentcore.ThinkingContent); ok { + if t.Redacted { + blocks = append(blocks, map[string]any{ + "type": "redacted_thinking", "data": t.ThinkingSignature, + }) + continue + } + if t.Thinking == "" { + continue + } + block := map[string]any{"type": "thinking", "thinking": t.Thinking} + if t.ThinkingSignature != "" { + block["signature"] = t.ThinkingSignature + } + blocks = append(blocks, block) + } + } + for _, c := range msg.Content { + switch b := c.(type) { + case agentcore.TextContent: + if b.Text == "" { + continue + } + blocks = append(blocks, map[string]any{"type": "text", "text": b.Text}) + case agentcore.ToolCallContent: + var input any + _ = json.Unmarshal(b.Arguments, &input) + blocks = append(blocks, map[string]any{ + "type": "tool_use", "id": b.ID, "name": b.Name, "input": input, + }) + } + } + if len(blocks) == 0 { + // No usable content: emit a single space rather than an empty text + // block ("text must be non-empty" is rejected by strict endpoints). + blocks = []map[string]any{{"type": "text", "text": " "}} + } + return map[string]any{"role": "assistant", "content": blocks} + case agentcore.ToolResultMessage: + return map[string]any{"role": "user", "content": []map[string]any{{ + "type": "tool_result", + "tool_use_id": msg.ToolCallID, + "content": agentcore.ContentToText(msg.Content), + }}} + default: + return nil + } +} + +// anthropicUserContent shapes a user content list for the Anthropic Messages +// wire. When there are no images it collapses to a plain string. When images +// are present it emits the content-block array form: text blocks plus image +// blocks with a base64 source ({"type":"image","source":{"type":"base64", +// "media_type":,"data":}}). +func anthropicUserContent(content agentcore.ContentList) any { + hasImage := false + for _, c := range content { + if _, ok := c.(agentcore.ImageContent); ok { + hasImage = true + break + } + } + if !hasImage { + return agentcore.ContentToText(content) + } + blocks := make([]map[string]any, 0, len(content)) + for _, c := range content { + switch b := c.(type) { + case agentcore.TextContent: + if b.Text == "" { + continue + } + blocks = append(blocks, map[string]any{"type": "text", "text": b.Text}) + case agentcore.ImageContent: + blocks = append(blocks, map[string]any{ + "type": "image", + "source": map[string]any{ + "type": "base64", + "media_type": b.MimeType, + "data": b.Data, + }, + }) + } + } + return blocks +} + +// contextHasImage reports whether any message in the context carries an image +// content block, so a driver can reject image input on a non-multimodal model. +func contextHasImage(msgs []agentcore.Message) bool { + for _, m := range msgs { + var content agentcore.ContentList + switch msg := m.(type) { + case agentcore.UserMessage: + content = msg.Content + case agentcore.CompactionMessage: + content = msg.AsUserMessage().Content + default: + continue + } + for _, c := range content { + if _, ok := c.(agentcore.ImageContent); ok { + return true + } + } + } + return false +} + +// checkImageSupport returns a clear error when the request carries image input +// but the named model (looked up in models) does not declare SupportsImages. +// A model absent from the catalog is treated permissively (unknown capability), +// deferring to the provider's own validation. This turns the silent drop of +// image blocks on a text-only model into an actionable message. +func checkImageSupport(providerName, model string, models []Model, msgs []agentcore.Message) error { + if !contextHasImage(msgs) { + return nil + } + for _, m := range models { + if m.ID == model { + if !m.SupportsImages { + return fmt.Errorf("%s: model %q does not support image input", providerName, model) + } + return nil + } + } + return nil +} + +// encodeAnthropicTools maps AgentTools onto the Anthropic tool schema. +func encodeAnthropicTools(tools []agentcore.AgentTool) []map[string]any { + if len(tools) == 0 { + return nil + } + out := make([]map[string]any, 0, len(tools)) + for _, t := range tools { + schema := json.RawMessage(t.Schema()) + if len(schema) == 0 { + schema = json.RawMessage("{}") + } + out = append(out, map[string]any{ + "name": t.Name(), + "description": t.Description(), + "input_schema": schema, + }) + } + return out +} + +// maxOutputTokensFor pulls a max-output hint from the request Extra map, if any, +// so callers can bound Anthropic responses without a registry lookup here. +func maxOutputTokensFor(req CompletionRequest) int { + if req.Config.Extra == nil { + return 0 + } + switch v := req.Config.Extra["max_tokens"].(type) { + case int: + return v + case float64: + return int(v) + } + return 0 +} + +// --------------------------------------------------------------------------- +// Constructors. +// --------------------------------------------------------------------------- + +// openAICompatPreset captures the per-gateway differences among the +// OpenAI-compatible providers: everything the constructors used to repeat +// (provider name, default endpoint, whether auth is required, and any extra +// headers). Collapsing the near-identical constructors onto this table means +// "add a gateway" is a one-line entry plus a thin exported wrapper. +type openAICompatPreset struct { + name string + defaultURL string // "" ⇒ no default; baseURL must be supplied by the caller + requiresAuth bool + extraHeaders map[string]string +} + +// newOpenAICompat builds an OpenAI-compatible driver from a preset, falling back +// to the preset's default endpoint when baseURL is empty. +func newOpenAICompat(p openAICompatPreset, baseURL string, models []Model) Provider { + if baseURL == "" { + baseURL = p.defaultURL + } + return &openAICompatDriver{ + name: p.name, + baseURL: baseURL, + models: models, + requiresAuth: p.requiresAuth, + extraHeaders: p.extraHeaders, + } +} + +// NewOpenRouterProvider builds the OpenRouter provider — the reference +// OpenAI-compatible gateway. baseURL defaults to the public endpoint when empty. +func NewOpenRouterProvider(baseURL string, models []Model) Provider { + return newOpenAICompat(openAICompatPreset{ + name: "openrouter", + defaultURL: openRouterBaseURL, + requiresAuth: true, + extraHeaders: map[string]string{ + // OpenRouter attribution headers (optional but recommended). + "HTTP-Referer": "https://github.com/smallnest/pigo", + "X-Title": "pigo", + }, + }, baseURL, models) +} + +// NewOllamaProvider builds the Ollama provider (local, OpenAI-compatible, no +// auth). baseURL defaults to the local daemon when empty. +func NewOllamaProvider(baseURL string, models []Model) Provider { + return newOpenAICompat(openAICompatPreset{ + name: "ollama", + defaultURL: ollamaBaseURL, + requiresAuth: false, + }, baseURL, models) +} + +// NewNvidiaProvider builds the NVIDIA provider (hosted NIM, OpenAI-compatible, +// Bearer auth). baseURL defaults to the public integrate endpoint when empty. +// The API key is resolved by the "nvidia" provider name (NVIDIA_API_KEY); +// secret values are never logged. +func NewNvidiaProvider(baseURL string, models []Model) Provider { + return newOpenAICompat(openAICompatPreset{ + name: "nvidia", + defaultURL: nvidiaBaseURL, + requiresAuth: true, + }, baseURL, models) +} + +// NewOpenAICompatibleProvider builds a generic OpenAI-compatible provider for an +// arbitrary gateway reached by baseURL (Bearer auth). It is the target of an +// explicit --protocol=openai selection: unlike the preset constructors it has no +// default endpoint (baseURL must be supplied) and carries the neutral provider +// name "openai", so an API key resolves from OPENAI_API_KEY (or the --api-key +// override bound to that name). Secret values are never logged. +func NewOpenAICompatibleProvider(baseURL string, models []Model) Provider { + return newOpenAICompat(openAICompatPreset{ + name: "openai", + defaultURL: "", // no default: caller must supply the endpoint + requiresAuth: true, + }, baseURL, models) +} + +// newAnthropicCompat builds an Anthropic-Messages driver, falling back to +// defaultURL when baseURL is empty. It is the shared body of the two +// Anthropic-wire constructors, which differ only in name, default endpoint, and +// auth header. +func newAnthropicCompat(name, defaultURL, baseURL string, models []Model, authHeader func(*http.Request, string)) Provider { + if baseURL == "" { + baseURL = defaultURL + } + return &anthropicCompatDriver{ + name: name, + baseURL: baseURL, + models: models, + authHeader: authHeader, + } +} + +// anthropicAuthHeaderFor returns the auth-header setter for an Anthropic-Messages +// provider given its registry AuthScheme (spec.AuthScheme). The two shapes seen +// among anthropic-protocol providers are: +// +// - AuthBearer → Authorization: Bearer . Used by anthropic-protocol +// gateways that authenticate with a plain bearer token on their /anthropic +// endpoint. +// - AuthXAPIKey → x-api-key: plus the required anthropic-version +// header. This is the direct-Anthropic convention and, per pi's behavior, +// also what MiniMax (minimax / minimax-cn) uses on its /anthropic endpoint. +// +// Any other scheme (e.g. AuthAWS for Bedrock, AuthSpecial) falls back to the +// x-api-key convention so a generic anthropic-protocol provider does not crash; +// bespoke auth for those is layered by a later node. The returned func never +// logs the secret value. +func anthropicAuthHeaderFor(authScheme string) func(*http.Request, string) { + if authScheme == AuthBearer { + return func(req *http.Request, apiKey string) { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + } + return func(req *http.Request, apiKey string) { + req.Header.Set("x-api-key", apiKey) + req.Header.Set("anthropic-version", anthropicAPIVersion) + } +} + +// NewAnthropicProvider builds a provider that speaks the Anthropic Messages wire +// format directly (POST {baseURL}/messages), the target of an explicit +// --protocol=anthropic selection. baseURL defaults to the public Anthropic API +// when empty. Auth uses the Anthropic conventions: an x-api-key header plus the +// required anthropic-version header. The API key resolves by the "anthropic" +// provider name (ANTHROPIC_API_KEY / CLAUDE_API_KEY, or the --api-key override); +// secret values are never logged. +func NewAnthropicProvider(baseURL string, models []Model) Provider { + return newAnthropicCompat("anthropic", anthropicBaseURL, baseURL, models, + anthropicAuthHeaderFor(AuthXAPIKey)) +} + +// NewAnthropicProtocolProvider builds a named Anthropic-Messages provider whose +// auth header follows the given registry AuthScheme. It is the target of an +// explicit --provider selection for any anthropic-protocol built-in (anthropic, +// minimax, minimax-cn, and — routed generically for now — bedrock/ +// cloudflare-ai-gateway): the driver identity is the provider's own name (so +// errors reference it), baseURL is the already-resolved endpoint (spec default +// or override), and authScheme selects the header shape (see +// anthropicAuthHeaderFor). Secret values are never logged. +func NewAnthropicProtocolProvider(name, baseURL, authScheme string, models []Model) Provider { + return newAnthropicCompat(name, anthropicBaseURL, baseURL, models, + anthropicAuthHeaderFor(authScheme)) +} + +// NewBedrockProvider builds the Bedrock provider, reusing the Anthropic Messages +// decoder (Anthropic-on-Bedrock speaks the Messages wire format). baseURL +// defaults to a us-east-1 runtime endpoint when empty; real deployments pass +// the region-specific URL. Auth is a Bearer token (Bedrock API keys); SigV4 +// signing, when required, is layered by the caller's HTTP client. +func NewBedrockProvider(baseURL string, models []Model) Provider { + return newAnthropicCompat("bedrock", bedrockBaseURL, baseURL, models, + func(req *http.Request, apiKey string) { + req.Header.Set("Authorization", "Bearer "+apiKey) + }) +} diff --git a/pigo/internal/provider/providers_anthropic_test.go b/pigo/internal/provider/providers_anthropic_test.go new file mode 100644 index 0000000..c9ba613 --- /dev/null +++ b/pigo/internal/provider/providers_anthropic_test.go @@ -0,0 +1,275 @@ +// Tests for the Anthropic-Messages-protocol built-in providers (US-006, node +// #187): anthropic, minimax, minimax-cn. They assert the registry metadata +// (base_url, protocol, key env var) and that the constructed anthropic-compat +// driver attaches the auth header dictated by the provider's AuthScheme. +// +// No real network calls are made: the auth header is exercised by invoking the +// driver's authHeader func against a dummy *http.Request and inspecting the +// resulting headers (the same package can reach the unexported driver fields). +package provider + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// TestAnthropicProtocolProviderRegistrySpecs asserts the registry metadata for +// each Anthropic-Messages-protocol built-in: default base URL, wire protocol, +// and the key env var (via envAPIKey + t.Setenv). +func TestAnthropicProtocolProviderRegistrySpecs(t *testing.T) { + cases := []struct { + name string + wantBaseURL string + wantEnv string + }{ + {"anthropic", "https://api.anthropic.com/v1", "ANTHROPIC_API_KEY"}, + {"minimax", "https://api.minimax.io/anthropic", "MINIMAX_API_KEY"}, + {"minimax-cn", "https://api.minimaxi.com/anthropic", "MINIMAX_CN_API_KEY"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + spec, ok := LookupProviderSpec(tc.name) + if !ok { + t.Fatalf("provider %q not found in registry", tc.name) + } + if spec.Protocol != ProtocolAnthropic { + t.Errorf("protocol = %q, want %q", spec.Protocol, ProtocolAnthropic) + } + if spec.DefaultBaseURL != tc.wantBaseURL { + t.Errorf("base_url = %q, want %q", spec.DefaultBaseURL, tc.wantBaseURL) + } + // Key resolution: the provider's key comes from tc.wantEnv. + t.Setenv(tc.wantEnv, "SEKRET-"+tc.name) + if got := envAPIKey(tc.name); got != "SEKRET-"+tc.name { + t.Errorf("envAPIKey(%q) = %q, want key resolved from %s", tc.name, got, tc.wantEnv) + } + }) + } +} + +// TestAnthropicProtocolProviderAuthHeader verifies that the driver built for +// each provider (the way resolveNamedProvider builds it: name + resolved +// base_url + spec.AuthScheme) targets the provider's base URL and sets the auth +// header matching its AuthScheme. anthropic/minimax/minimax-cn are all +// x-api-key + anthropic-version per pi's convention. +func TestAnthropicProtocolProviderAuthHeader(t *testing.T) { + for _, name := range []string{"anthropic", "minimax", "minimax-cn"} { + t.Run(name, func(t *testing.T) { + spec, ok := LookupProviderSpec(name) + if !ok { + t.Fatalf("provider %q not found in registry", name) + } + p := NewAnthropicProtocolProvider(spec.Name, spec.DefaultBaseURL, spec.AuthScheme, nil) + d, ok := p.(*anthropicCompatDriver) + if !ok { + t.Fatalf("provider %q is not an *anthropicCompatDriver", name) + } + if d.baseURL != spec.DefaultBaseURL { + t.Errorf("baseURL = %q, want %q", d.baseURL, spec.DefaultBaseURL) + } + req, err := http.NewRequest(http.MethodPost, d.baseURL+"/messages", nil) + if err != nil { + t.Fatalf("build request: %v", err) + } + d.authHeader(req, "SEKRET") + assertAuthHeaderForScheme(t, req, spec.AuthScheme) + if got := req.Header.Get("Authorization"); got != "" && spec.AuthScheme != AuthBearer { + t.Errorf("unexpected Authorization header %q for x-api-key scheme", got) + } + }) + } +} + +// TestAnthropicAuthSchemeSelection proves the auth-header mechanism itself: +// AuthBearer yields Authorization: Bearer, AuthXAPIKey yields x-api-key plus the +// anthropic-version header. This guarantees an anthropic-protocol provider can +// select either header shape from its registry AuthScheme. +func TestAnthropicAuthSchemeSelection(t *testing.T) { + t.Run("bearer", func(t *testing.T) { + p := NewAnthropicProtocolProvider("gw", "https://example.test/anthropic", AuthBearer, nil) + req := newDummyReq(t) + p.(*anthropicCompatDriver).authHeader(req, "SEKRET") + if got := req.Header.Get("Authorization"); got != "Bearer SEKRET" { + t.Errorf("Authorization = %q, want %q", got, "Bearer SEKRET") + } + if got := req.Header.Get("x-api-key"); got != "" { + t.Errorf("x-api-key = %q, want empty for bearer scheme", got) + } + }) + t.Run("x-api-key", func(t *testing.T) { + p := NewAnthropicProtocolProvider("anthropic", "", AuthXAPIKey, nil) + req := newDummyReq(t) + p.(*anthropicCompatDriver).authHeader(req, "SEKRET") + if got := req.Header.Get("x-api-key"); got != "SEKRET" { + t.Errorf("x-api-key = %q, want %q", got, "SEKRET") + } + if got := req.Header.Get("anthropic-version"); got != anthropicAPIVersion { + t.Errorf("anthropic-version = %q, want %q", got, anthropicAPIVersion) + } + }) + // A non-crashing fallback for unwired schemes (e.g. Bedrock's AuthAWS): must + // not panic and defaults to the x-api-key convention. + t.Run("fallback", func(t *testing.T) { + p := NewAnthropicProtocolProvider("bedrock", "", AuthAWS, nil) + req := newDummyReq(t) + p.(*anthropicCompatDriver).authHeader(req, "SEKRET") + if got := req.Header.Get("x-api-key"); got != "SEKRET" { + t.Errorf("fallback x-api-key = %q, want %q", got, "SEKRET") + } + }) +} + +// TestNewAnthropicProviderUnchanged guards against regressing the direct +// Anthropic constructor: it must still default to the public endpoint and use +// x-api-key + anthropic-version. +func TestNewAnthropicProviderUnchanged(t *testing.T) { + d, ok := NewAnthropicProvider("", nil).(*anthropicCompatDriver) + if !ok { + t.Fatal("NewAnthropicProvider did not return *anthropicCompatDriver") + } + if d.baseURL != anthropicBaseURL { + t.Errorf("baseURL = %q, want %q", d.baseURL, anthropicBaseURL) + } + req := newDummyReq(t) + d.authHeader(req, "SEKRET") + if got := req.Header.Get("x-api-key"); got != "SEKRET" { + t.Errorf("x-api-key = %q, want %q", got, "SEKRET") + } + if got := req.Header.Get("anthropic-version"); got != anthropicAPIVersion { + t.Errorf("anthropic-version = %q, want %q", got, anthropicAPIVersion) + } +} + +func newDummyReq(t *testing.T) *http.Request { + t.Helper() + req, err := http.NewRequest(http.MethodPost, "https://example.test/messages", nil) + if err != nil { + t.Fatalf("build request: %v", err) + } + return req +} + +func assertAuthHeaderForScheme(t *testing.T, req *http.Request, scheme string) { + t.Helper() + if scheme == AuthBearer { + if got := req.Header.Get("Authorization"); got != "Bearer SEKRET" { + t.Errorf("Authorization = %q, want %q", got, "Bearer SEKRET") + } + return + } + if got := req.Header.Get("x-api-key"); got != "SEKRET" { + t.Errorf("x-api-key = %q, want %q", got, "SEKRET") + } + if got := req.Header.Get("anthropic-version"); got != anthropicAPIVersion { + t.Errorf("anthropic-version = %q, want %q", got, anthropicAPIVersion) + } +} + +// TestEncodeAnthropicRequestIncludesModel guards a required field of the +// Anthropic Messages API: the request body must carry the "model" id. Omitting +// it makes the public API return 400 and OpenAI-compatible gateways return an +// empty/error response, which the SSE decoder silently turns into an empty +// assistant turn — a confusing "no output, no error" failure. +func TestEncodeAnthropicRequestIncludesModel(t *testing.T) { + req := CompletionRequest{ + Model: "claude-opus-4-8", + } + body, err := encodeAnthropicRequest(req, nil) + if err != nil { + t.Fatalf("encodeAnthropicRequest: %v", err) + } + var decoded map[string]any + if err := json.Unmarshal(body, &decoded); err != nil { + t.Fatalf("unmarshal body: %v", err) + } + if got, ok := decoded["model"]; !ok { + t.Fatalf("request body missing required \"model\" field; keys=%v", keysOf(decoded)) + } else if got != "claude-opus-4-8" { + t.Errorf("model = %v, want claude-opus-4-8", got) + } +} + +func keysOf(m map[string]any) []string { + ks := make([]string, 0, len(m)) + for k := range m { + ks = append(ks, k) + } + return ks +} + +// TestEncodeAnthropicRequestCoalescesToolResults guards the Anthropic protocol +// rule that every tool_use must be answered by tool_result blocks in a single +// immediately-following user message. Before the fix, each tool result was +// encoded as its own user turn, which the API rejected with "tool_use ids were +// found without tool_result blocks immediately after" whenever one assistant +// turn carried multiple tool_use calls. +func TestEncodeAnthropicRequestCoalescesToolResults(t *testing.T) { + req := CompletionRequest{ + Model: "claude-x", + Context: LlmContext{ + Messages: agentcore.MessageList{ + // Assistant turn with two tool_use blocks. + agentcore.AssistantMessage{ + Content: agentcore.ContentList{ + agentcore.NewToolCallContent("call_1", "read", json.RawMessage(`{}`)), + agentcore.NewToolCallContent("call_2", "write", json.RawMessage(`{}`)), + }, + }, + // Two consecutive tool results for that turn. + agentcore.ToolResultMessage{ + ToolCallID: "call_1", + Content: agentcore.ContentList{agentcore.NewTextContent("result-1")}, + }, + agentcore.ToolResultMessage{ + ToolCallID: "call_2", + Content: agentcore.ContentList{agentcore.NewTextContent("result-2")}, + }, + // Next assistant turn (text reply). + agentcore.AssistantMessage{ + Content: agentcore.ContentList{agentcore.NewTextContent("done")}, + }, + }, + }, + } + + body, err := encodeAnthropicRequest(req, nil) + if err != nil { + t.Fatalf("encodeAnthropicRequest: %v", err) + } + decoded := decodeBody(t, body) + msgs, ok := decoded["messages"].([]any) + if !ok { + t.Fatalf("messages not an array: %T", decoded["messages"]) + } + if len(msgs) != 3 { + t.Fatalf("want 3 top-level messages (assistant tool_use, one user tool_result turn, assistant reply), got %d:\n%v", len(msgs), msgs) + } + + // Message 2 must be the single user turn holding BOTH tool_result blocks. + second, ok := msgs[1].(map[string]any) + if !ok { + t.Fatalf("message[1] not an object: %T", msgs[1]) + } + if second["role"] != "user" { + t.Fatalf("message[1].role = %v, want user", second["role"]) + } + blocks, ok := second["content"].([]any) + if !ok { + t.Fatalf("message[1].content not an array: %T", second["content"]) + } + if len(blocks) != 2 { + t.Fatalf("want 2 tool_result blocks in the coalesced turn, got %d:\n%v", len(blocks), second["content"]) + } + for i, wantID := range []string{"call_1", "call_2"} { + block, ok := blocks[i].(map[string]any) + if !ok { + t.Fatalf("block[%d] not an object: %T", i, blocks[i]) + } + if block["type"] != "tool_result" || block["tool_use_id"] != wantID { + t.Errorf("block[%d] = %v, want tool_result for %s", i, block, wantID) + } + } +} diff --git a/pigo/internal/provider/providers_openai_test.go b/pigo/internal/provider/providers_openai_test.go new file mode 100644 index 0000000..4e3db7e --- /dev/null +++ b/pigo/internal/provider/providers_openai_test.go @@ -0,0 +1,151 @@ +// Node #186: end-to-end wiring test for every OpenAI-protocol built-in provider. +// +// For each provider reachable via --provider (US-005) this asserts three things: +// (a) the registry spec reports Protocol "openai" and the PRD-mandated default +// base URL, +// (b) its primary API-key env var resolves through envAPIKey (the same path +// auth.go uses at request time), and +// (c) the generic OpenAI-compatible construction path used by main.go's +// resolveNamedProvider builds a non-nil driver bound to the spec's model. +// +// It also pins the restored legacy key aliases (CLAUDE_API_KEY, GOOGLE_API_KEY, +// NVIDIA_NIM_API_KEY) and the non-standard HF_TOKEN key name. +// +// This file is intentionally separate from providers_test.go: a sibling node +// edits that file concurrently. +package provider + +import "testing" + +// openAIWiringCase describes one OpenAI-protocol provider's expected registry +// metadata and primary key env var. +type openAIWiringCase struct { + name string + baseURL string + primaryEnv string +} + +// openAIProviderCases lists every OpenAI-protocol provider that must work +// end-to-end via --provider (US-005). Base URLs mirror the PRD's Technical +// Considerations table. +var openAIProviderCases = []openAIWiringCase{ + {"groq", "https://api.groq.com/openai/v1", "GROQ_API_KEY"}, + {"xai", "https://api.x.ai/v1", "XAI_API_KEY"}, + {"cerebras", "https://api.cerebras.ai/v1", "CEREBRAS_API_KEY"}, + {"mistral", "https://api.mistral.ai", "MISTRAL_API_KEY"}, + {"moonshotai", "https://api.moonshot.ai/v1", "MOONSHOT_API_KEY"}, + {"moonshotai-cn", "https://api.moonshot.cn/v1", "MOONSHOT_API_KEY"}, + {"fireworks", "https://api.fireworks.ai/inference", "FIREWORKS_API_KEY"}, + {"together", "https://api.together.ai/v1", "TOGETHER_API_KEY"}, + {"openrouter", "https://openrouter.ai/api/v1", "OPENROUTER_API_KEY"}, + {"nvidia", "https://integrate.api.nvidia.com/v1", "NVIDIA_API_KEY"}, + {"zai", "https://api.z.ai/api/coding/paas/v4", "ZAI_API_KEY"}, + {"zai-coding-cn", "https://open.bigmodel.cn/api/coding/paas/v4", "ZAI_CODING_CN_API_KEY"}, + {"kimi-coding", "https://api.kimi.com/coding", "KIMI_API_KEY"}, + {"opencode", "https://opencode.ai/zen", "OPENCODE_API_KEY"}, + {"opencode-go", "https://opencode.ai/zen/go", "OPENCODE_API_KEY"}, + {"huggingface", "https://router.huggingface.co/v1", "HF_TOKEN"}, + {"ant-ling", "https://api.ant-ling.com/v1", "ANT_LING_API_KEY"}, + {"vercel-ai-gateway", "https://ai-gateway.vercel.sh", "AI_GATEWAY_API_KEY"}, + {"xiaomi", "https://api.xiaomimimo.com/v1", "XIAOMI_API_KEY"}, + {"xiaomi-token-plan-cn", "https://token-plan-cn.xiaomimimo.com/v1", "XIAOMI_TOKEN_PLAN_CN_API_KEY"}, + {"xiaomi-token-plan-ams", "https://token-plan-ams.xiaomimimo.com/v1", "XIAOMI_TOKEN_PLAN_AMS_API_KEY"}, + {"xiaomi-token-plan-sgp", "https://token-plan-sgp.xiaomimimo.com/v1", "XIAOMI_TOKEN_PLAN_SGP_API_KEY"}, +} + +func TestOpenAIProviderWiring(t *testing.T) { + for _, tc := range openAIProviderCases { + t.Run(tc.name, func(t *testing.T) { + // (a) registry spec: protocol + default base URL + primary env var. + spec, ok := LookupProviderSpec(tc.name) + if !ok { + t.Fatalf("LookupProviderSpec(%q): not found in registry", tc.name) + } + if spec.Protocol != ProtocolOpenAI { + t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolOpenAI) + } + if spec.DefaultBaseURL != tc.baseURL { + t.Errorf("DefaultBaseURL = %q, want %q", spec.DefaultBaseURL, tc.baseURL) + } + if len(spec.EnvVars) == 0 || spec.EnvVars[0] != tc.primaryEnv { + t.Errorf("primary EnvVar = %v, want first = %q", spec.EnvVars, tc.primaryEnv) + } + + // (b) key resolution via the primary env var, using the same + // envAPIKey path auth.go relies on at request time. + t.Setenv(tc.primaryEnv, "sk-"+tc.name) + if got := envAPIKey(tc.name); got != "sk-"+tc.name { + t.Errorf("envAPIKey(%q) = %q, want %q", tc.name, got, "sk-"+tc.name) + } + + // (c) construction path equivalent to main.go's resolveNamedProvider + // for an openai-protocol spec: build a generic OpenAI-compatible + // driver against the spec's base URL, bound to the spec's model. + models := []Model{{Provider: spec.Name, ID: "test-model", SupportsImages: true}} + drv := NewOpenAICompatibleProvider(spec.DefaultBaseURL, models) + if drv == nil { + t.Fatalf("NewOpenAICompatibleProvider(%q) returned nil", spec.DefaultBaseURL) + } + got := drv.Models() + if len(got) != 1 || got[0].Provider != spec.Name || got[0].ID != "test-model" { + t.Errorf("driver Models() = %+v, want one model bound to provider %q", got, spec.Name) + } + }) + } +} + +// TestLegacyKeyAliases pins the secondary env vars restored to the registry so +// credentials set under older names still resolve. +func TestLegacyKeyAliases(t *testing.T) { + aliases := []struct { + provider string + primary string // primary env var (must NOT be set for the alias to be exercised) + alias string // legacy alias env var under test + }{ + {"anthropic", "ANTHROPIC_API_KEY", "CLAUDE_API_KEY"}, + {"google", "GEMINI_API_KEY", "GOOGLE_API_KEY"}, + {"nvidia", "NVIDIA_API_KEY", "NVIDIA_NIM_API_KEY"}, + } + for _, a := range aliases { + t.Run(a.provider, func(t *testing.T) { + // Ensure the alias is listed in the registry's EnvVars. + spec, ok := LookupProviderSpec(a.provider) + if !ok { + t.Fatalf("LookupProviderSpec(%q): not found", a.provider) + } + found := false + for _, e := range spec.EnvVars { + if e == a.alias { + found = true + break + } + } + if !found { + t.Fatalf("EnvVars %v missing legacy alias %q", spec.EnvVars, a.alias) + } + // Clear the primary so only the alias can satisfy resolution, then + // assert the alias resolves. + t.Setenv(a.primary, "") + t.Setenv(a.alias, "legacy-"+a.provider) + if got := envAPIKey(a.provider); got != "legacy-"+a.provider { + t.Errorf("envAPIKey(%q) via %s = %q, want %q", a.provider, a.alias, got, "legacy-"+a.provider) + } + }) + } +} + +// TestHuggingFaceTokenEnv asserts the non-standard HF_TOKEN key name resolves +// for huggingface (it does not follow the _API_KEY convention). +func TestHuggingFaceTokenEnv(t *testing.T) { + spec, ok := LookupProviderSpec("huggingface") + if !ok { + t.Fatal("LookupProviderSpec(\"huggingface\"): not found") + } + if len(spec.EnvVars) == 0 || spec.EnvVars[0] != "HF_TOKEN" { + t.Fatalf("huggingface EnvVars = %v, want first = HF_TOKEN", spec.EnvVars) + } + t.Setenv("HF_TOKEN", "hf-secret") + if got := envAPIKey("huggingface"); got != "hf-secret" { + t.Errorf("envAPIKey(\"huggingface\") = %q, want %q", got, "hf-secret") + } +} diff --git a/pigo/internal/provider/providers_test.go b/pigo/internal/provider/providers_test.go new file mode 100644 index 0000000..42c794d --- /dev/null +++ b/pigo/internal/provider/providers_test.go @@ -0,0 +1,227 @@ +package provider + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// captureServer records the last request path, headers, and decoded JSON body, +// then replays a canned SSE stream. It lets the provider tests assert the wire +// shape a driver produced without a live upstream. +type captureServer struct { + srv *httptest.Server + path string + headers http.Header + body map[string]any +} + +func newCaptureServer(t *testing.T, sseBody string) *captureServer { + t.Helper() + cs := &captureServer{} + cs.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cs.path = r.URL.Path + cs.headers = r.Header.Clone() + raw, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(raw, &cs.body) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(sseBody)) + })) + t.Cleanup(cs.srv.Close) + return cs +} + +// drainStream collects event kinds and the final message from a provider stream. +func drainStream(t *testing.T, stream *AssistantMessageEventStream) ([]string, agentcore.AssistantMessage) { + t.Helper() + var kinds []string + for ev := range stream.Events() { + kinds = append(kinds, ev.EventKind()) + } + final, err := stream.Result(context.Background()) + if err != nil { + t.Fatalf("stream result: %v", err) + } + return kinds, final +} + +// TestOpenRouterProviderStreamsChatCompletions drives OpenRouter (the reference +// OpenAI-compatible provider) end to end: the driver must POST to +// /chat/completions with a Bearer token and stream through OpenAIDecoder. +func TestOpenRouterProviderStreamsChatCompletions(t *testing.T) { + cs := newCaptureServer(t, openaiToolCallSSE) + p := NewOpenRouterProvider(cs.srv.URL, []Model{{Provider: "openrouter", ID: "openai/gpt-4o"}}) + + if p.Name() != "openrouter" { + t.Errorf("name = %q, want openrouter", p.Name()) + } + stream, err := p.StreamCompletion(context.Background(), CompletionRequest{ + Model: "openai/gpt-4o", + Context: LlmContext{SystemPrompt: "be brief", Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion: %v", err) + } + kinds, final := drainStream(t, stream) + if kinds[len(kinds)-1] != StreamEventDone { + t.Errorf("last event = %q, want done", kinds[len(kinds)-1]) + } + if final.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason = %q, want tool_use", final.StopReason) + } + + // Wire assertions: path, auth header, and request body shape. + if cs.path != "/chat/completions" { + t.Errorf("path = %q, want /chat/completions", cs.path) + } + if got := cs.headers.Get("Authorization"); got != "Bearer sk-test" { + t.Errorf("auth header = %q, want Bearer sk-test", got) + } + if cs.body["stream"] != true { + t.Errorf("stream flag = %v, want true", cs.body["stream"]) + } + if cs.body["model"] != "openai/gpt-4o" { + t.Errorf("model = %v, want openai/gpt-4o", cs.body["model"]) + } + msgs, _ := cs.body["messages"].([]any) + if len(msgs) != 2 { + t.Fatalf("messages len = %d, want 2 (system+user)", len(msgs)) + } + first, _ := msgs[0].(map[string]any) + if first["role"] != "system" || first["content"] != "be brief" { + t.Errorf("system message = %v", first) + } +} + +// TestOpenRouterMissingKeyIsEarlyError verifies a missing API key is the early +// "cannot build the stream" returned error, and the provider name is named +// without leaking any value. +func TestOpenRouterMissingKeyIsEarlyError(t *testing.T) { + p := NewOpenRouterProvider("", nil) + _, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"}) + if err == nil { + t.Fatal("missing API key must return an early error") + } + if !strings.Contains(err.Error(), "openrouter") { + t.Errorf("error should name the provider, got %v", err) + } +} + +// TestOllamaProviderNoAuth verifies the local Ollama provider sends no +// Authorization header and still streams through OpenAIDecoder. +func TestOllamaProviderNoAuth(t *testing.T) { + cs := newCaptureServer(t, openaiToolCallSSE) + p := NewOllamaProvider(cs.srv.URL, []Model{{Provider: "ollama", ID: "llama3"}}) + if p.Name() != "ollama" { + t.Errorf("name = %q, want ollama", p.Name()) + } + // No API key configured — Ollama must not require one. + stream, err := p.StreamCompletion(context.Background(), CompletionRequest{ + Model: "llama3", + Context: LlmContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}}, + }) + if err != nil { + t.Fatalf("StreamCompletion (no auth): %v", err) + } + _, final := drainStream(t, stream) + if final.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason = %q, want tool_use", final.StopReason) + } + if got := cs.headers.Get("Authorization"); got != "" { + t.Errorf("Ollama must send no auth header, got %q", got) + } +} + +// TestBedrockProviderStreamsAnthropicWire verifies the Bedrock provider POSTs +// the Anthropic Messages wire format (system top-level, max_tokens present) and +// streams through AnthropicDecoder. +func TestBedrockProviderStreamsAnthropicWire(t *testing.T) { + cs := newCaptureServer(t, anthropicToolUseSSE) + p := NewBedrockProvider(cs.srv.URL, []Model{{Provider: "bedrock", ID: "anthropic.claude-3"}}) + if p.Name() != "bedrock" { + t.Errorf("name = %q, want bedrock", p.Name()) + } + stream, err := p.StreamCompletion(context.Background(), CompletionRequest{ + Model: "anthropic.claude-3", + Context: LlmContext{SystemPrompt: "sys", Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}}, + Config: StreamConfig{APIKey: "bedrock-key"}, + }) + if err != nil { + t.Fatalf("StreamCompletion: %v", err) + } + kinds, _ := drainStream(t, stream) + if kinds[len(kinds)-1] != StreamEventDone { + t.Errorf("last event = %q, want done", kinds[len(kinds)-1]) + } + if cs.body["system"] != "sys" { + t.Errorf("system = %v, want top-level 'sys'", cs.body["system"]) + } + if cs.body["max_tokens"] == nil { + t.Error("Anthropic wire requires max_tokens") + } + if cs.body["stream"] != true { + t.Errorf("stream flag = %v, want true", cs.body["stream"]) + } +} + +// TestBedrockMissingKeyIsEarlyError mirrors the OpenRouter early-error check. +func TestBedrockMissingKeyIsEarlyError(t *testing.T) { + p := NewBedrockProvider("", nil) + _, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"}) + if err == nil { + t.Fatal("missing API key must return an early error") + } + if !strings.Contains(err.Error(), "bedrock") { + t.Errorf("error should name the provider, got %v", err) + } +} + +// TestNvidiaProviderStreamsChatCompletions verifies the NVIDIA NIM provider +// rides the OpenAI-compatible driver: it POSTs to /chat/completions with a +// Bearer token and streams through OpenAIDecoder, exactly like OpenRouter. +func TestNvidiaProviderStreamsChatCompletions(t *testing.T) { + cs := newCaptureServer(t, openaiToolCallSSE) + p := NewNvidiaProvider(cs.srv.URL, []Model{{Provider: "nvidia", ID: "meta/llama-3.3-70b-instruct"}}) + if p.Name() != "nvidia" { + t.Errorf("name = %q, want nvidia", p.Name()) + } + stream, err := p.StreamCompletion(context.Background(), CompletionRequest{ + Model: "meta/llama-3.3-70b-instruct", + Context: LlmContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}}, + Config: StreamConfig{APIKey: "nvapi-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion: %v", err) + } + kinds, _ := drainStream(t, stream) + if kinds[len(kinds)-1] != StreamEventDone { + t.Errorf("last event = %q, want done", kinds[len(kinds)-1]) + } + if cs.path != "/chat/completions" { + t.Errorf("path = %q, want /chat/completions", cs.path) + } + if got := cs.headers.Get("Authorization"); got != "Bearer nvapi-test" { + t.Errorf("auth header = %q, want Bearer nvapi-test", got) + } +} + +// TestNvidiaMissingKeyIsEarlyError verifies NVIDIA requires an API key and +// reports the missing key as an early error naming the provider (never a value). +func TestNvidiaMissingKeyIsEarlyError(t *testing.T) { + p := NewNvidiaProvider("", nil) + _, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"}) + if err == nil { + t.Fatal("missing API key must return an early error") + } + if !strings.Contains(err.Error(), "nvidia") { + t.Errorf("error should name the provider, got %v", err) + } +} diff --git a/pigo/internal/provider/registry.go b/pigo/internal/provider/registry.go new file mode 100644 index 0000000..3abfcf6 --- /dev/null +++ b/pigo/internal/provider/registry.go @@ -0,0 +1,377 @@ +// This file defines the central provider registry (US-001): a single source of +// truth for every built-in provider's metadata — its name, the environment +// variables that carry its API key (in precedence order), its default base URL, +// wire protocol, auth scheme, any extra headers, and provider-specific base-URL +// override env vars. +// +// The registry is deliberately additive and self-contained: later nodes wire +// auth resolution (auth.go), the --provider flag (main.go), base_url overrides, +// and per-provider construction (providers.go) to READ from it. This node only +// introduces the data + a lookup, so it does not change existing behavior. +// +// Data source: the PRD "Technical Considerations" table +// (tasks/prd-provider-env-parity.md), derived from pi's env-api-keys.ts and the +// per-provider *.models.ts files. +// +// Security: the registry holds only env var NAMES, never secret values. Keys are +// resolved from the environment at request time (see auth.go) and never logged. +package provider + +// ProviderSpec is the metadata describing one built-in provider. It is the +// single source of truth consumed by auth resolution, the --provider flag, +// base_url override handling, and per-provider wiring. +type ProviderSpec struct { + // Name is the canonical provider name (e.g. "deepseek", "zai-coding-cn"). + Name string + // EnvVars lists the environment variables checked (in precedence order) for + // this provider's API key. The first non-empty value wins. + EnvVars []string + // DefaultBaseURL is the provider's default API endpoint. It may be a template + // (containing placeholders like {region}) for providers whose endpoint is + // composed from additional parameters (Bedrock, Vertex, Cloudflare, Azure). + DefaultBaseURL string + // Protocol is the wire protocol the provider speaks: "openai" (OpenAI Chat + // Completions) or "anthropic" (Anthropic Messages). + Protocol string + // AuthScheme names how credentials are attached: "bearer", "x-api-key", + // "aws", "azure", or "special". + AuthScheme string + // ExtraHeaders are provider-specific headers attached to every request (may + // be nil). + ExtraHeaders map[string]string + // BaseURLEnvVars lists provider-specific base-URL override environment + // variables (e.g. AZURE_OPENAI_BASE_URL), in precedence order. May be empty; + // the generic _BASE_URL convention is handled by callers. + BaseURLEnvVars []string +} + +// Protocol values. +const ( + ProtocolOpenAI = "openai" + ProtocolAnthropic = "anthropic" +) + +// AuthScheme values. +const ( + AuthBearer = "bearer" + AuthXAPIKey = "x-api-key" + AuthAWS = "aws" + AuthAzure = "azure" + AuthSpecial = "special" +) + +// providerRegistry is the ordered list of all built-in provider specs. Order is +// stable so callers that enumerate providers (e.g. --help) get a deterministic +// list. LookupProviderSpec indexes it by name. +var providerRegistry = []ProviderSpec{ + { + Name: "anthropic", + EnvVars: []string{"ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY", "CLAUDE_API_KEY"}, + DefaultBaseURL: anthropicBaseURL, // https://api.anthropic.com/v1 + Protocol: ProtocolAnthropic, + AuthScheme: AuthXAPIKey, + }, + { + Name: "openai", + EnvVars: []string{"OPENAI_API_KEY"}, + DefaultBaseURL: "https://api.openai.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "ant-ling", + EnvVars: []string{"ANT_LING_API_KEY"}, + DefaultBaseURL: "https://api.ant-ling.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "deepseek", + EnvVars: []string{"DEEPSEEK_API_KEY"}, + DefaultBaseURL: "https://api.deepseek.com", + Protocol: ProtocolOpenAIResponses, + AuthScheme: AuthBearer, + }, + { + Name: "nvidia", + EnvVars: []string{"NVIDIA_API_KEY", "NVIDIA_NIM_API_KEY"}, + DefaultBaseURL: nvidiaBaseURL, // https://integrate.api.nvidia.com/v1 + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "google", + EnvVars: []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"}, + DefaultBaseURL: "https://generativelanguage.googleapis.com/v1beta", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "groq", + EnvVars: []string{"GROQ_API_KEY"}, + DefaultBaseURL: "https://api.groq.com/openai/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "cerebras", + EnvVars: []string{"CEREBRAS_API_KEY"}, + DefaultBaseURL: "https://api.cerebras.ai/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "xai", + EnvVars: []string{"XAI_API_KEY"}, + DefaultBaseURL: "https://api.x.ai/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "openrouter", + EnvVars: []string{"OPENROUTER_API_KEY"}, + DefaultBaseURL: openRouterBaseURL, // https://openrouter.ai/api/v1 + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "vercel-ai-gateway", + EnvVars: []string{"AI_GATEWAY_API_KEY"}, + DefaultBaseURL: "https://ai-gateway.vercel.sh", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "zai", + EnvVars: []string{"ZAI_API_KEY"}, + DefaultBaseURL: "https://api.z.ai/api/coding/paas/v4", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "zai-coding-cn", + EnvVars: []string{"ZAI_CODING_CN_API_KEY"}, + DefaultBaseURL: "https://open.bigmodel.cn/api/coding/paas/v4", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "mistral", + EnvVars: []string{"MISTRAL_API_KEY"}, + DefaultBaseURL: "https://api.mistral.ai", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "minimax", + EnvVars: []string{"MINIMAX_API_KEY"}, + DefaultBaseURL: "https://api.minimax.io/anthropic", + Protocol: ProtocolAnthropic, + AuthScheme: AuthXAPIKey, + }, + { + Name: "minimax-cn", + EnvVars: []string{"MINIMAX_CN_API_KEY"}, + DefaultBaseURL: "https://api.minimaxi.com/anthropic", + Protocol: ProtocolAnthropic, + AuthScheme: AuthXAPIKey, + }, + { + Name: "moonshotai", + EnvVars: []string{"MOONSHOT_API_KEY"}, + DefaultBaseURL: "https://api.moonshot.ai/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "moonshotai-cn", + EnvVars: []string{"MOONSHOT_API_KEY"}, + DefaultBaseURL: "https://api.moonshot.cn/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "huggingface", + EnvVars: []string{"HF_TOKEN"}, + DefaultBaseURL: "https://router.huggingface.co/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "fireworks", + EnvVars: []string{"FIREWORKS_API_KEY"}, + DefaultBaseURL: "https://api.fireworks.ai/inference", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "together", + EnvVars: []string{"TOGETHER_API_KEY"}, + DefaultBaseURL: "https://api.together.ai/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "opencode", + EnvVars: []string{"OPENCODE_API_KEY"}, + DefaultBaseURL: "https://opencode.ai/zen", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "opencode-go", + EnvVars: []string{"OPENCODE_API_KEY"}, + DefaultBaseURL: "https://opencode.ai/zen/go", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "kimi-coding", + EnvVars: []string{"KIMI_API_KEY"}, + DefaultBaseURL: "https://api.kimi.com/coding", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "xiaomi", + EnvVars: []string{"XIAOMI_API_KEY"}, + DefaultBaseURL: "https://api.xiaomimimo.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "xiaomi-token-plan-cn", + EnvVars: []string{"XIAOMI_TOKEN_PLAN_CN_API_KEY"}, + DefaultBaseURL: "https://token-plan-cn.xiaomimimo.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "xiaomi-token-plan-ams", + EnvVars: []string{"XIAOMI_TOKEN_PLAN_AMS_API_KEY"}, + DefaultBaseURL: "https://token-plan-ams.xiaomimimo.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "xiaomi-token-plan-sgp", + EnvVars: []string{"XIAOMI_TOKEN_PLAN_SGP_API_KEY"}, + DefaultBaseURL: "https://token-plan-sgp.xiaomimimo.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + // Chinese cloud LLM platforms. All four expose OpenAI-compatible endpoints + // authenticated with a plain Bearer API key, so they reuse the standard + // OpenAI-compatible driver with no bespoke auth. Base URLs are the platforms' + // OpenAI-compatible endpoints as documented at implementation time. + { + // Baidu AI Cloud Qianfan (Baidu Qianfan). + Name: "qianfan", + EnvVars: []string{"QIANFAN_API_KEY"}, + DefaultBaseURL: "https://qianfan.baidubce.com/v2", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + // ByteDance Volcengine Ark (Volcengine Ark). ARK_API_KEY is the platform's + // conventional variable; VOLCENGINE_API_KEY is accepted as a fallback. + Name: "volcengine", + EnvVars: []string{"ARK_API_KEY", "VOLCENGINE_API_KEY"}, + DefaultBaseURL: "https://ark.cn-beijing.volces.com/api/v3", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + // Alibaba Cloud DashScope (DashScope), OpenAI-compatible mode. + Name: "dashscope", + EnvVars: []string{"DASHSCOPE_API_KEY"}, + DefaultBaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + // Tencent Hunyuan (Hunyuan), OpenAI-compatible endpoint. + Name: "hunyuan", + EnvVars: []string{"HUNYUAN_API_KEY"}, + DefaultBaseURL: "https://api.hunyuan.cloud.tencent.com/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "azure-openai-responses", + EnvVars: []string{"AZURE_OPENAI_API_KEY"}, + // Endpoint is composed from AZURE_OPENAI_BASE_URL / AZURE_OPENAI_RESOURCE_NAME + // per the Azure OpenAI convention; there is no fixed public default. + DefaultBaseURL: "", + Protocol: ProtocolOpenAI, + AuthScheme: AuthAzure, + BaseURLEnvVars: []string{"AZURE_OPENAI_BASE_URL"}, + }, + { + Name: "amazon-bedrock", + EnvVars: []string{"AWS_BEARER_TOKEN_BEDROCK"}, + // Region-specific runtime endpoint; {AWS_REGION} defaults to us-east-1. + DefaultBaseURL: "https://bedrock-runtime.{AWS_REGION}.amazonaws.com", + Protocol: ProtocolAnthropic, + AuthScheme: AuthAWS, + }, + { + Name: "google-vertex", + EnvVars: []string{"GOOGLE_CLOUD_API_KEY"}, + // Location-specific endpoint; protocol varies by model (Gemini vs Claude). + DefaultBaseURL: "https://{location}-aiplatform.googleapis.com", + Protocol: ProtocolOpenAI, + AuthScheme: AuthSpecial, + }, + { + Name: "cloudflare-workers-ai", + EnvVars: []string{"CLOUDFLARE_API_KEY"}, + // {id} is CLOUDFLARE_ACCOUNT_ID. + DefaultBaseURL: "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1", + Protocol: ProtocolOpenAI, + AuthScheme: AuthBearer, + }, + { + Name: "cloudflare-ai-gateway", + EnvVars: []string{"CLOUDFLARE_API_KEY"}, + // {acct}/{gw} are CLOUDFLARE_ACCOUNT_ID / CLOUDFLARE_GATEWAY_ID. + DefaultBaseURL: "https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/anthropic", + Protocol: ProtocolAnthropic, + AuthScheme: AuthXAPIKey, + }, +} + +// providerRegistryByName indexes providerRegistry by provider name for O(1) +// lookup. Built once at package init. +var providerRegistryByName = func() map[string]ProviderSpec { + m := make(map[string]ProviderSpec, len(providerRegistry)) + for _, spec := range providerRegistry { + m[spec.Name] = spec + } + return m +}() + +// LookupProviderSpec returns the ProviderSpec for a provider name and whether it +// is a known built-in provider. The returned spec is a copy; mutating its slice +// or map fields is discouraged as they are shared with the registry. +func LookupProviderSpec(name string) (ProviderSpec, bool) { + spec, ok := providerRegistryByName[name] + return spec, ok +} + +// ProviderSpecs returns all built-in provider specs in registry (display) order. +// Callers must not mutate the returned specs' slice/map fields. +func ProviderSpecs() []ProviderSpec { + out := make([]ProviderSpec, len(providerRegistry)) + copy(out, providerRegistry) + return out +} + +// ProviderNames returns all built-in provider names in registry order. +func ProviderNames() []string { + out := make([]string, len(providerRegistry)) + for i, spec := range providerRegistry { + out[i] = spec.Name + } + return out +} diff --git a/pigo/internal/provider/registry_test.go b/pigo/internal/provider/registry_test.go new file mode 100644 index 0000000..def540b --- /dev/null +++ b/pigo/internal/provider/registry_test.go @@ -0,0 +1,164 @@ +package provider + +import ( + "sort" + "testing" +) + +func TestLookupProviderSpec_Hit(t *testing.T) { + spec, ok := LookupProviderSpec("deepseek") + if !ok { + t.Fatalf("LookupProviderSpec(deepseek): expected hit, got miss") + } + if spec.Name != "deepseek" { + t.Errorf("Name = %q, want deepseek", spec.Name) + } + if spec.DefaultBaseURL != "https://api.deepseek.com" { + t.Errorf("DefaultBaseURL = %q, want https://api.deepseek.com", spec.DefaultBaseURL) + } + if spec.Protocol != ProtocolOpenAIResponses { + t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolOpenAIResponses) + } + if len(spec.EnvVars) != 1 || spec.EnvVars[0] != "DEEPSEEK_API_KEY" { + t.Errorf("EnvVars = %v, want [DEEPSEEK_API_KEY]", spec.EnvVars) + } +} + +func TestLookupProviderSpec_Miss(t *testing.T) { + if _, ok := LookupProviderSpec("does-not-exist"); ok { + t.Errorf("LookupProviderSpec(does-not-exist): expected miss, got hit") + } +} + +func TestAnthropicEnvVarOrder(t *testing.T) { + spec, ok := LookupProviderSpec("anthropic") + if !ok { + t.Fatal("LookupProviderSpec(anthropic): expected hit") + } + want := []string{"ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY", "CLAUDE_API_KEY"} + if len(spec.EnvVars) != len(want) { + t.Fatalf("EnvVars = %v, want %v", spec.EnvVars, want) + } + for i := range want { + if spec.EnvVars[i] != want[i] { + t.Errorf("EnvVars[%d] = %q, want %q (OAuth must be first)", i, spec.EnvVars[i], want[i]) + } + } + if spec.Protocol != ProtocolAnthropic { + t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolAnthropic) + } +} + +func TestHuggingfaceEnvVar(t *testing.T) { + spec, ok := LookupProviderSpec("huggingface") + if !ok { + t.Fatal("LookupProviderSpec(huggingface): expected hit") + } + if len(spec.EnvVars) != 1 || spec.EnvVars[0] != "HF_TOKEN" { + t.Errorf("EnvVars = %v, want [HF_TOKEN]", spec.EnvVars) + } +} + +func TestChineseCloudProviders(t *testing.T) { + cases := []struct { + name string + envVars []string + baseURL string + }{ + {"qianfan", []string{"QIANFAN_API_KEY"}, "https://qianfan.baidubce.com/v2"}, + {"volcengine", []string{"ARK_API_KEY", "VOLCENGINE_API_KEY"}, "https://ark.cn-beijing.volces.com/api/v3"}, + {"dashscope", []string{"DASHSCOPE_API_KEY"}, "https://dashscope.aliyuncs.com/compatible-mode/v1"}, + {"hunyuan", []string{"HUNYUAN_API_KEY"}, "https://api.hunyuan.cloud.tencent.com/v1"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + spec, ok := LookupProviderSpec(tc.name) + if !ok { + t.Fatalf("LookupProviderSpec(%q): expected hit", tc.name) + } + if spec.Name != tc.name { + t.Errorf("Name = %q, want %q", spec.Name, tc.name) + } + if spec.DefaultBaseURL != tc.baseURL { + t.Errorf("DefaultBaseURL = %q, want %q", spec.DefaultBaseURL, tc.baseURL) + } + if spec.Protocol != ProtocolOpenAI { + t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolOpenAI) + } + if spec.AuthScheme != AuthBearer { + t.Errorf("AuthScheme = %q, want %q", spec.AuthScheme, AuthBearer) + } + if len(spec.EnvVars) != len(tc.envVars) { + t.Fatalf("EnvVars = %v, want %v", spec.EnvVars, tc.envVars) + } + for i := range tc.envVars { + if spec.EnvVars[i] != tc.envVars[i] { + t.Errorf("EnvVars[%d] = %q, want %q", i, spec.EnvVars[i], tc.envVars[i]) + } + } + }) + } +} + +func TestAzureBaseURLEnvVars(t *testing.T) { + spec, ok := LookupProviderSpec("azure-openai-responses") + if !ok { + t.Fatal("LookupProviderSpec(azure-openai-responses): expected hit") + } + found := false + for _, v := range spec.BaseURLEnvVars { + if v == "AZURE_OPENAI_BASE_URL" { + found = true + } + } + if !found { + t.Errorf("BaseURLEnvVars = %v, want to contain AZURE_OPENAI_BASE_URL", spec.BaseURLEnvVars) + } +} + +func TestRegistryContainsAllExpectedProviders(t *testing.T) { + expected := []string{ + "anthropic", "openai", "ant-ling", "deepseek", "nvidia", "google", + "groq", "cerebras", "xai", "openrouter", "vercel-ai-gateway", "zai", + "zai-coding-cn", "mistral", "minimax", "minimax-cn", "moonshotai", + "moonshotai-cn", "huggingface", "fireworks", "together", "opencode", + "opencode-go", "kimi-coding", "xiaomi", "xiaomi-token-plan-cn", + "xiaomi-token-plan-ams", "xiaomi-token-plan-sgp", + "qianfan", "volcengine", "dashscope", "hunyuan", + "azure-openai-responses", "amazon-bedrock", "google-vertex", + "cloudflare-workers-ai", "cloudflare-ai-gateway", + } + for _, name := range expected { + if _, ok := LookupProviderSpec(name); !ok { + t.Errorf("registry missing expected provider %q", name) + } + } + names := ProviderNames() + if len(names) != len(expected) { + t.Errorf("registry has %d providers, want %d", len(names), len(expected)) + } + + // Every spec must have a name, at least one env var, and a valid protocol. + for _, spec := range ProviderSpecs() { + if spec.Name == "" { + t.Error("found spec with empty Name") + } + if len(spec.EnvVars) == 0 { + t.Errorf("provider %q has no EnvVars", spec.Name) + } + if spec.Protocol != ProtocolOpenAI && spec.Protocol != ProtocolAnthropic && spec.Protocol != ProtocolOpenAIResponses { + t.Errorf("provider %q has invalid Protocol %q", spec.Name, spec.Protocol) + } + } + + // No duplicate provider names. + seen := make(map[string]bool, len(names)) + dupSorted := append([]string(nil), names...) + sort.Strings(dupSorted) + for _, n := range dupSorted { + if seen[n] { + t.Errorf("duplicate provider name %q in registry", n) + } + seen[n] = true + } +} diff --git a/pigo/internal/provider/resolve.go b/pigo/internal/provider/resolve.go new file mode 100644 index 0000000..82f9547 --- /dev/null +++ b/pigo/internal/provider/resolve.go @@ -0,0 +1,223 @@ +package provider + +// Provider resolution moved here from cmd/pigo (US-004, #361): mapping a model +// id / --provider / --protocol selection to a concrete wire driver, plus the +// base-url override precedence. Environment lookups are injected as an +// env func(string) string so callers (and tests) control the environment +// instead of reaching into the process env directly. + +import ( + "fmt" + "strings" + + "github.com/smallnest/pigo/internal/cli/config" +) + +// ResolveProvider maps a model id to a built-in provider. An explicit +// --provider name wins over every other rule: it selects a built-in provider +// from the registry and constructs the matching wire driver (see +// ResolveNamedProvider). When provider is empty, protocol and model-id +// heuristics apply as before. +// +// When protocol is a non-empty explicit selection ("openai" or "anthropic") it +// wins over the model-id heuristics: the provider is built directly for that +// wire format against baseURL, which is how a user points pigo at a self-hosted +// or third-party endpoint and says which protocol it speaks. An "anthropic" +// selection with no baseURL targets the public Anthropic API. +// +// When protocol is empty, resolution falls back to model-id heuristics: +// +// 1. If the id is in the preset catalog, use its declared provider (this is how +// OpenRouter/NVIDIA/Ollama presets pick the right gateway). +// 2. An "ollama/" prefix (or a base URL on the Ollama port) → local Ollama. +// 3. An "nvidia/" prefix → NVIDIA NIM (strips the prefix for the wire id). +// 4. Model-name inference: with no --base-url, a well-known model-name prefix +// (e.g. "claude-*", "deepseek-*") selects its first-party built-in provider +// via ResolveNamedProvider (see InferProviderFromModel). +// 5. Everything else → OpenRouter, the reference OpenAI-compatible gateway. +// +// An unknown protocol value is an error, surfaced to the caller for exit-code +// mapping rather than silently falling back. +func ResolveProvider(model, baseURL, protocol, providerName string, env func(string) string) (Provider, string, error) { + // Explicit --provider selects a built-in provider from the registry and + // wins over both --protocol inference and model-id heuristics. + if strings.TrimSpace(providerName) != "" { + return ResolveNamedProvider(providerName, model, baseURL, protocol, env) + } + + // 0. Explicit protocol selection wins over every heuristic. Normalize the + // surface value first so "openai" and "openai/chat" collapse to the same + // Chat Completions selector and "openai/resp_api" routes to the Responses + // driver; an unknown value surfaces as an error for exit-code mapping. + canonical, err := NormalizeProtocol(protocol) + if err != nil { + return nil, "", err + } + switch canonical { + case ProtocolOpenAI: + if strings.TrimSpace(baseURL) == "" { + return nil, "", fmt.Errorf("--protocol openai requires --base-url") + } + return NewOpenAICompatibleProvider(baseURL, []Model{{Provider: "openai", ID: model, SupportsImages: true}}), "openai", nil + case ProtocolOpenAIResponses: + // The Responses driver has no public default endpoint here: unlike the + // anthropic path (which targets the public API), resp_api mirrors the + // Chat Completions requirement and demands an explicit --base-url. + if strings.TrimSpace(baseURL) == "" { + return nil, "", fmt.Errorf("--protocol openai/resp_api requires --base-url") + } + return NewOpenAIResponsesProvider("openai", baseURL, []Model{{Provider: "openai", ID: model, SupportsImages: true}}), "openai", nil + case ProtocolAnthropic: + return NewAnthropicProvider(baseURL, []Model{{Provider: "anthropic", ID: model, SupportsImages: true}}), "anthropic", nil + case "": + // fall through to heuristic resolution + } + + // 1. Preset catalog wins: a curated id knows its own provider. + if p, ok := LookupPreset(model); ok { + switch p.Provider { + case "nvidia": + return NewNvidiaProvider(baseURL, []Model{{Provider: "nvidia", ID: model, SupportsImages: true}}), "nvidia", nil + case "ollama": + id := strings.TrimPrefix(model, "ollama/") + return NewOllamaProvider(baseURL, []Model{{Provider: "ollama", ID: id, SupportsImages: true}}), "ollama", nil + case "", "openrouter": + return NewOpenRouterProvider(baseURL, []Model{{Provider: "openrouter", ID: model, SupportsImages: true}}), "openrouter", nil + default: + // Any other preset provider is a named built-in (e.g. deepseek, + // qianfan, dashscope): build it from the registry so the correct + // base URL, protocol, and API-key env var are used — not OpenRouter's. + return ResolveNamedProvider(p.Provider, model, baseURL, protocol, env) + } + } + + // 2. Local Ollama by prefix or port. + if strings.HasPrefix(model, "ollama/") || strings.Contains(baseURL, "11434") { + id := strings.TrimPrefix(model, "ollama/") + return NewOllamaProvider(baseURL, []Model{{Provider: "ollama", ID: id, SupportsImages: true}}), "ollama", nil + } + // 3. NVIDIA NIM by prefix. + if strings.HasPrefix(model, "nvidia/") { + id := strings.TrimPrefix(model, "nvidia/") + return NewNvidiaProvider(baseURL, []Model{{Provider: "nvidia", ID: id, SupportsImages: true}}), "nvidia", nil + } + // 4. Model-name inference: with no --provider/--protocol (both empty here) and + // no --base-url, guess the provider from the model name's well-known prefix + // (e.g. "claude-*" → anthropic, "deepseek-*" → deepseek). A confident hit is + // routed through ResolveNamedProvider so the provider's registry protocol, + // default base URL, and API-key env var are used. A --base-url is treated as + // a custom-endpoint signal that should not be second-guessed, so inference is + // skipped when one is given. Ambiguous/unknown names fall through to (5). + if strings.TrimSpace(baseURL) == "" { + if name, ok := InferProviderFromModel(model); ok { + return ResolveNamedProvider(name, model, baseURL, protocol, env) + } + } + // 5. Default: OpenRouter. + return NewOpenRouterProvider(baseURL, []Model{{Provider: "openrouter", ID: model, SupportsImages: true}}), "openrouter", nil +} + +// ResolveNamedProvider builds the driver for an explicit --provider selection. +// It looks the name up in the built-in registry and constructs the wire driver +// matching the spec's Protocol: "openai" → an OpenAI-compatible (Bearer) driver, +// "anthropic" → an Anthropic-Messages driver. The base URL follows the override +// precedence in ResolveBaseURL (--base-url > provider-specific env > generic +// _BASE_URL > spec default). The returned provider-name string is the +// spec name, so downstream API-key resolution reads the provider's own env var +// (spec.EnvVars). +// +// Special providers with bespoke auth (azure/bedrock/vertex/cloudflare — +// AuthScheme aws/azure/special, or the cloudflare-* names) are routed to +// ResolveSpecialProvider, which validates their required env vars and composes +// the concrete endpoint (node #188). +func ResolveNamedProvider(name, model, baseURL, protocol string, env func(string) string) (Provider, string, error) { + spec, ok := LookupProviderSpec(name) + if !ok { + return nil, "", fmt.Errorf("unknown --provider %q (available: %s)", name, strings.Join(ProviderNames(), ", ")) + } + // A concurrently-set --protocol must agree with the provider's own protocol; + // an incompatible pair is a user error naming both flags. Normalize the raw + // value first so aliases (e.g. "openai/chat" for an "openai" spec) don't + // falsely conflict, and a genuine typo surfaces as a clear "unknown --protocol" + // error rather than a misleading conflict message. + if strings.TrimSpace(protocol) != "" { + canonical, err := NormalizeProtocol(protocol) + if err != nil { + return nil, "", err + } + if canonical != spec.Protocol { + return nil, "", fmt.Errorf("--provider %q speaks the %q protocol, which conflicts with --protocol %q; drop --protocol or set it to %q", name, spec.Protocol, protocol, spec.Protocol) + } + } + // Special-auth providers (Azure / Bedrock / Vertex / Cloudflare) compose + // their endpoint from several env vars and/or need non-standard credential + // validation, so route them to the dedicated resolver (US-007 / node #188). + // It performs its own base-URL composition (honoring the --base-url override) + // and returns a clear error naming any absent required env var. + if IsSpecialAuthProvider(spec) { + p, err := ResolveSpecialProvider(spec, model, baseURL, env) + if err != nil { + return nil, "", err + } + return p, spec.Name, nil + } + // Base-URL precedence (US-004 / FR-8, FR-9): --base-url flag > provider- + // specific base-url env var(s) > generic _BASE_URL > spec default. + url := ResolveBaseURL(spec, baseURL, env) + models := []Model{{Provider: spec.Name, ID: model, SupportsImages: true}} + // Note: spec.ExtraHeaders would be attached here, but the exported generic + // constructors do not yet accept custom headers; all built-in specs currently + // carry no ExtraHeaders, so this is a no-op today (refined alongside #188). + switch spec.Protocol { + case ProtocolAnthropic: + // Auth header follows the spec's AuthScheme (x-api-key + anthropic-version + // for anthropic/minimax/minimax-cn; Bearer for any anthropic-protocol + // gateway that authenticates with a plain bearer token). The driver name is + // the spec name so errors reference the selected provider. + return NewAnthropicProtocolProvider(spec.Name, url, spec.AuthScheme, models), spec.Name, nil + case ProtocolOpenAI: + return NewOpenAICompatibleProvider(url, models), spec.Name, nil + case ProtocolOpenAIResponses: + return NewOpenAIResponsesProvider(spec.Name, url, models), spec.Name, nil + default: + // The registry only ever stores openai/openai-resp/anthropic; guard anyway + // so an unexpected value is a clear error rather than a nil provider. + return nil, "", fmt.Errorf("--provider %q has unsupported protocol %q", name, spec.Protocol) + } +} + +// ResolveBaseURL determines the effective base URL for a selected provider, +// applying the base_url override precedence (US-004 / FR-8, FR-9). The first +// non-empty source wins, in this order: +// +// 1. flagBaseURL — the explicit --base-url/-u flag (highest). +// 2. provider-specific base-url env var(s) from spec.BaseURLEnvVars, in the +// order the registry declares them (e.g. AZURE_OPENAI_BASE_URL). +// 3. the generic _BASE_URL env var, where is the provider +// name uppercased with '-' rewritten to '_' (e.g. zai-coding-cn → +// ZAI_CODING_CN_BASE_URL). +// 4. spec.DefaultBaseURL — the registry default (lowest). +// +// Values are trimmed of surrounding whitespace before the non-empty check, so a +// whitespace-only env var does not shadow a lower-precedence source. Environment +// lookups go through the injected env func so callers control the environment. +func ResolveBaseURL(spec ProviderSpec, flagBaseURL string, env func(string) string) string { + // 1. Explicit flag wins over every env-var convention. + if v := strings.TrimSpace(flagBaseURL); v != "" { + return v + } + // 2. Provider-specific override env vars, in registry precedence order. + for _, name := range spec.BaseURLEnvVars { + if v := strings.TrimSpace(env(name)); v != "" { + return v + } + } + // 3. Generic _BASE_URL convention. + if envName := config.GenericBaseURLEnvVar(spec.Name); envName != "" { + if v := strings.TrimSpace(env(envName)); v != "" { + return v + } + } + // 4. Registry default. + return spec.DefaultBaseURL +} diff --git a/pigo/internal/provider/resolve_test.go b/pigo/internal/provider/resolve_test.go new file mode 100644 index 0000000..3771f92 --- /dev/null +++ b/pigo/internal/provider/resolve_test.go @@ -0,0 +1,296 @@ +package provider + +// Tests for provider resolution moved from cmd/pigo (US-004, #361): ResolveProvider +// maps a model id to the right gateway (preset catalog first, then prefix rules, +// then OpenRouter default), and ResolveBaseURL applies the base-url override +// precedence. Environment lookups are injected via os.Getenv here. + +import ( + "os" + "strings" + "testing" +) + +// TestResolveProviderPresetCatalog verifies a preset id resolves to its declared +// provider (NVIDIA and Ollama presets do not fall through to OpenRouter). +func TestResolveProviderPresetCatalog(t *testing.T) { + cases := []struct { + model string + wantName string + }{ + {"meta/llama-3.3-70b-instruct", "nvidia"}, // NVIDIA preset + {"ollama/llama3.3", "ollama"}, // Ollama preset + {"openai/gpt-4o", "openrouter"}, // OpenRouter preset + {"anthropic/claude-3.5-sonnet", "openrouter"}, // OpenRouter preset + } + for _, c := range cases { + _, name, err := ResolveProvider(c.model, "", "", "", os.Getenv) + if err != nil { + t.Errorf("ResolveProvider(%q) error: %v", c.model, err) + continue + } + if name != c.wantName { + t.Errorf("ResolveProvider(%q) = %q, want %q", c.model, name, c.wantName) + } + } +} + +// TestResolveProviderPrefixAndDefault verifies the prefix rules and the +// OpenRouter default for ids not in the catalog. +func TestResolveProviderPrefixAndDefault(t *testing.T) { + cases := []struct { + model string + baseURL string + wantName string + }{ + {"ollama/some-local-model", "", "ollama"}, // ollama/ prefix + {"nvidia/some-nim-model", "", "nvidia"}, // nvidia/ prefix + {"some-unknown-model", "", "openrouter"}, // default + {"m", "http://host:11434/v1", "ollama"}, // ollama port + } + for _, c := range cases { + _, name, err := ResolveProvider(c.model, c.baseURL, "", "", os.Getenv) + if err != nil { + t.Errorf("ResolveProvider(%q) error: %v", c.model, err) + continue + } + if name != c.wantName { + t.Errorf("ResolveProvider(%q, %q) = %q, want %q", c.model, c.baseURL, name, c.wantName) + } + } +} + +// TestResolveProviderExplicitProtocol verifies an explicit --protocol wins over +// model-id heuristics: openai (with base-url) and anthropic select the matching +// wire driver, an empty base-url for openai errors, and an unknown protocol +// errors instead of silently falling back. +func TestResolveProviderExplicitProtocol(t *testing.T) { + // openai protocol → "openai" provider name, requires base-url. + if _, name, err := ResolveProvider("any-model", "https://example.com/v1", "openai", "", os.Getenv); err != nil || name != "openai" { + t.Errorf("protocol=openai = (%q, %v), want (openai, nil)", name, err) + } + if _, _, err := ResolveProvider("any-model", "", "openai", "", os.Getenv); err == nil { + t.Error("protocol=openai with no base-url should error") + } + // anthropic protocol → "anthropic" provider name, base-url optional (defaults). + if _, name, err := ResolveProvider("claude-x", "", "anthropic", "", os.Getenv); err != nil || name != "anthropic" { + t.Errorf("protocol=anthropic = (%q, %v), want (anthropic, nil)", name, err) + } + // Unknown protocol errors rather than falling back to a heuristic. + if _, _, err := ResolveProvider("any-model", "", "grpc", "", os.Getenv); err == nil { + t.Error("unknown protocol should error") + } +} + +// TestResolveProviderResponsesProtocol verifies the openai/resp_api selector +// routes to the Responses driver (against an explicit base-url), that the +// "openai/chat" alias resolves identically to "openai", and that resp_api with +// no base-url errors like the plain openai path (mirroring the base-url +// requirement rather than defaulting to a public endpoint). +func TestResolveProviderResponsesProtocol(t *testing.T) { + // openai/resp_api → "openai" provider name, backed by the Responses driver. + p, name, err := ResolveProvider("any-model", "https://example.com/v1", "openai/resp_api", "", os.Getenv) + if err != nil || name != "openai" { + t.Fatalf("protocol=openai/resp_api = (%q, %v), want (openai, nil)", name, err) + } + if _, ok := p.(*responsesDriver); !ok { + t.Errorf("protocol=openai/resp_api built %T, want *responsesDriver", p) + } + // resp_api with no base-url errors, mirroring the openai requirement. + if _, _, err := ResolveProvider("any-model", "", "openai/resp_api", "", os.Getenv); err == nil { + t.Error("protocol=openai/resp_api with no base-url should error") + } + // "openai/chat" is an alias of "openai": same driver, same base-url rule. + p, name, err = ResolveProvider("any-model", "https://example.com/v1", "openai/chat", "", os.Getenv) + if err != nil || name != "openai" { + t.Fatalf("protocol=openai/chat = (%q, %v), want (openai, nil)", name, err) + } + if _, ok := p.(*responsesDriver); ok { + t.Error("protocol=openai/chat should build the Chat Completions driver, not *responsesDriver") + } + if _, _, err := ResolveProvider("any-model", "", "openai/chat", "", os.Getenv); err == nil { + t.Error("protocol=openai/chat with no base-url should error") + } +} + +// TestResolveProviderExplicitProvider verifies that --provider selects a +// built-in provider from the registry: the returned provider-name is the spec +// name (so key resolution reads the right env var), an OpenAI-protocol provider +// (deepseek) and an Anthropic-protocol provider (minimax) both resolve, an +// incompatible --protocol is a conflict error naming both flags, and an unknown +// provider name errors while listing the available names. +func TestResolveProviderExplicitProvider(t *testing.T) { + // OpenAI-protocol provider: returns its own name for key lookup. + if _, name, err := ResolveProvider("deepseek-chat", "", "", "deepseek", os.Getenv); err != nil || name != "deepseek" { + t.Errorf("provider=deepseek = (%q, %v), want (deepseek, nil)", name, err) + } + // Anthropic-protocol provider. + if _, name, err := ResolveProvider("MiniMax-M2", "", "", "minimax", os.Getenv); err != nil || name != "minimax" { + t.Errorf("provider=minimax = (%q, %v), want (minimax, nil)", name, err) + } + // A matching --protocol is not a conflict (deepseek speaks openai/resp_api). + if _, name, err := ResolveProvider("deepseek-chat", "", "openai/resp_api", "deepseek", os.Getenv); err != nil || name != "deepseek" { + t.Errorf("provider=deepseek + protocol=openai/resp_api = (%q, %v), want (deepseek, nil)", name, err) + } + // --provider wins over model-id heuristics: an ollama/-prefixed id still + // resolves to the named provider, not local Ollama. + if _, name, err := ResolveProvider("ollama/x", "", "", "deepseek", os.Getenv); err != nil || name != "deepseek" { + t.Errorf("provider=deepseek with ollama/ model = (%q, %v), want (deepseek, nil)", name, err) + } + // --base-url overrides the spec default without changing the provider name. + if _, name, err := ResolveProvider("deepseek-chat", "https://proxy.local/v1", "", "deepseek", os.Getenv); err != nil || name != "deepseek" { + t.Errorf("provider=deepseek + base-url = (%q, %v), want (deepseek, nil)", name, err) + } + // Conflict: minimax speaks anthropic; forcing --protocol openai errors and + // names both flags. + _, _, err := ResolveProvider("MiniMax-M2", "", "openai", "minimax", os.Getenv) + if err == nil { + t.Fatal("provider=minimax + protocol=openai should conflict") + } + if !strings.Contains(err.Error(), "--provider") || !strings.Contains(err.Error(), "--protocol") { + t.Errorf("conflict error should name both flags, got: %v", err) + } + // Unknown provider errors and lists available names. + _, _, err = ResolveProvider("m", "", "", "no-such-provider", os.Getenv) + if err == nil { + t.Fatal("unknown provider should error") + } + if !strings.Contains(err.Error(), "deepseek") { + t.Errorf("unknown-provider error should list available names, got: %v", err) + } + // An invalid --protocol paired with a named provider surfaces the clear + // "unknown --protocol" error (listing the accepted set) rather than a + // misleading conflict message. + _, _, err = ResolveProvider("deepseek-chat", "", "openai_api", "deepseek", os.Getenv) + if err == nil { + t.Fatal("provider=deepseek + protocol=openai_api should error") + } + if !strings.Contains(err.Error(), "unknown --protocol") { + t.Errorf("invalid --protocol should surface the unknown-protocol error, got: %v", err) + } +} + +// TestResolveProviderCNPresets verifies the Chinese-cloud preset ids route to +// their own provider (not the OpenRouter default) via the LookupPreset branch. +func TestResolveProviderCNPresets(t *testing.T) { + cases := []struct { + model string + wantName string + }{ + {"ernie-4.5-turbo-32k", "qianfan"}, + {"doubao-seed-1-6", "volcengine"}, + {"qwen-max", "dashscope"}, + {"hunyuan-turbos-latest", "hunyuan"}, + } + for _, c := range cases { + _, name, err := ResolveProvider(c.model, "", "", "", os.Getenv) + if err != nil { + t.Errorf("ResolveProvider(%q) error: %v", c.model, err) + continue + } + if name != c.wantName { + t.Errorf("ResolveProvider(%q) = %q, want %q", c.model, name, c.wantName) + } + } +} + +// TestResolveProviderCNExplicit verifies --provider selects the CN providers +// directly and that --base-url overrides without changing the provider name. +func TestResolveProviderCNExplicit(t *testing.T) { + for _, name := range []string{"qianfan", "volcengine", "dashscope", "hunyuan"} { + if _, got, err := ResolveProvider("some-model", "", "", name, os.Getenv); err != nil || got != name { + t.Errorf("provider=%s = (%q, %v), want (%s, nil)", name, got, err, name) + } + if _, got, err := ResolveProvider("some-model", "https://proxy.local/v1", "", name, os.Getenv); err != nil || got != name { + t.Errorf("provider=%s + base-url = (%q, %v), want (%s, nil)", name, got, err, name) + } + } +} + +// TestResolveProviderModelNameInference verifies model-name inference (Issue +// #235): with only --model given, a bare model name whose prefix identifies a +// single provider resolves to that provider — NOT the OpenRouter default. +func TestResolveProviderModelNameInference(t *testing.T) { + cases := []struct { + model string + wantName string + }{ + {"claude-opus-4-8", "anthropic"}, + {"deepseek-chat", "deepseek"}, + {"gpt-4.1", "openai"}, + {"gemini-3-pro", "google"}, + {"grok-5", "xai"}, + } + for _, c := range cases { + if _, name, err := ResolveProvider(c.model, "", "", "", os.Getenv); err != nil || name != c.wantName { + t.Errorf("ResolveProvider(%q) = (%q, %v), want (%q, nil)", c.model, name, err, c.wantName) + } + } +} + +// TestResolveProviderInferencePrecedence verifies that model-name inference does +// not override explicit flags and does not fire when a --base-url is given, and +// that unknown/ambiguous names still fall back to OpenRouter. +func TestResolveProviderInferencePrecedence(t *testing.T) { + // Explicit --provider wins over an inferable model name. + if _, name, err := ResolveProvider("claude-opus-4-8", "", "", "deepseek", os.Getenv); err != nil || name != "deepseek" { + t.Errorf("provider=deepseek overrides inference = (%q, %v), want (deepseek, nil)", name, err) + } + // Explicit --protocol wins over an inferable model name. + if _, name, err := ResolveProvider("claude-opus-4-8", "https://example.com/v1", "openai", "", os.Getenv); err != nil || name != "openai" { + t.Errorf("protocol=openai overrides inference = (%q, %v), want (openai, nil)", name, err) + } + // A --base-url signals a custom endpoint: inference is skipped, default applies. + if _, name, err := ResolveProvider("claude-opus-4-8", "https://gw.local/v1", "", "", os.Getenv); err != nil || name != "openrouter" { + t.Errorf("inference skipped with base-url = (%q, %v), want (openrouter, nil)", name, err) + } + // Ambiguous/unknown names still default to OpenRouter. + for _, m := range []string{"llama-3.3-70b", "totally-unknown-model"} { + if _, name, err := ResolveProvider(m, "", "", "", os.Getenv); err != nil || name != "openrouter" { + t.Errorf("ResolveProvider(%q) = (%q, %v), want (openrouter, nil)", m, name, err) + } + } +} + +// TestResolveBaseURLPrecedence exercises all four precedence levels for a +// hyphenated provider (zai-coding-cn → ZAI_CODING_CN_BASE_URL). +func TestResolveBaseURLPrecedence(t *testing.T) { + spec, ok := LookupProviderSpec("zai-coding-cn") + if !ok { + t.Fatal("expected zai-coding-cn in registry") + } + if got := ResolveBaseURL(spec, "", os.Getenv); got != spec.DefaultBaseURL { + t.Errorf("default: got %q, want %q", got, spec.DefaultBaseURL) + } + t.Setenv("ZAI_CODING_CN_BASE_URL", "https://generic.example/v4") + if got := ResolveBaseURL(spec, "", os.Getenv); got != "https://generic.example/v4" { + t.Errorf("generic env: got %q, want %q", got, "https://generic.example/v4") + } + if got := ResolveBaseURL(spec, "https://flag.example/v4", os.Getenv); got != "https://flag.example/v4" { + t.Errorf("flag over generic: got %q, want %q", got, "https://flag.example/v4") + } +} + +// TestResolveBaseURLProviderSpecificEnv covers a provider that declares a +// provider-specific base-url env var (azure), asserting it sits between the flag +// and the generic convention in precedence. +func TestResolveBaseURLProviderSpecificEnv(t *testing.T) { + spec, ok := LookupProviderSpec("azure-openai-responses") + if !ok { + t.Fatal("expected azure-openai-responses in registry") + } + if len(spec.BaseURLEnvVars) == 0 { + t.Fatal("expected azure-openai-responses to declare BaseURLEnvVars") + } + t.Setenv("AZURE_OPENAI_BASE_URL", "https://specific.example") + if got := ResolveBaseURL(spec, "", os.Getenv); got != "https://specific.example" { + t.Errorf("provider-specific env: got %q, want %q", got, "https://specific.example") + } + t.Setenv("AZURE_OPENAI_RESPONSES_BASE_URL", "https://generic.example") + if got := ResolveBaseURL(spec, "", os.Getenv); got != "https://specific.example" { + t.Errorf("provider-specific beats generic: got %q, want %q", got, "https://specific.example") + } + if got := ResolveBaseURL(spec, "https://flag.example", os.Getenv); got != "https://flag.example" { + t.Errorf("flag beats provider-specific: got %q, want %q", got, "https://flag.example") + } +} diff --git a/pigo/internal/provider/responses.go b/pigo/internal/provider/responses.go new file mode 100644 index 0000000..694bc4c --- /dev/null +++ b/pigo/internal/provider/responses.go @@ -0,0 +1,474 @@ +// This file implements the OpenAI Responses API driver (US-003, #539), the +// backing for --protocol openai/resp_api. Unlike openAICompatDriver (which +// hand-rolls the Chat Completions wire format), this driver speaks the +// Responses API (POST {base_url}/responses) via the official +// github.com/openai/openai-go SDK. +// +// This milestone covers streaming text: a plain prompt in, assistant text out, +// consumed from the Responses SSE stream and mapped into pigo's AssistantMessage +// the same way OpenAIDecoder does (API/Provider tags, Usage, +// ResponseID/ResponseModel, StopReason=end_turn). Tools (#541) and +// images/reasoning (#542) layer on later. +// +// Failure model (FR-13): only the earliest "cannot build the stream" case +// (missing API key) is a returned error. Every runtime failure — including a +// non-2xx from the endpoint — rides the returned stream as a terminal +// StreamErrorEvent, matching the chat driver's observable behavior. +// +// Base URL + auth: the SDK client is pointed at the resolved base_url and given +// the resolved key via option.WithBaseURL / option.WithAPIKey rather than +// reading the environment, so ResolveBaseURL precedence is preserved. A custom +// *http.Client may be injected (option.WithHTTPClient) for tests. +package provider + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + "github.com/openai/openai-go" + "github.com/openai/openai-go/option" + "github.com/openai/openai-go/responses" + "github.com/openai/openai-go/shared" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// responsesDriver is the Provider backing --protocol openai/resp_api. It holds +// the provider identity, the resolved endpoint, the model catalog, and whether +// an API key is required, and builds an openai-go client per request. +type responsesDriver struct { + name string + baseURL string + models []Model + // requiresAuth reports whether an API key must be present. Public OpenAI / + // Azure require it; a local gateway may not. + requiresAuth bool + // clientOpts are extra SDK options; tests inject option.WithHTTPClient here + // to stub the transport. + clientOpts []option.RequestOption +} + +// NewOpenAIResponsesProvider builds a Responses API provider targeting baseURL. +// baseURL must be the fully resolved endpoint (e.g. https://api.openai.com/v1); +// the SDK appends the /responses path. +func NewOpenAIResponsesProvider(name, baseURL string, models []Model) *responsesDriver { + return &responsesDriver{ + name: name, + baseURL: baseURL, + models: models, + requiresAuth: true, + } +} + +func (d *responsesDriver) Name() string { return d.name } +func (d *responsesDriver) Models() []Model { return d.models } + +// StreamCompletion issues a streaming Responses API call and surfaces the result +// on an AssistantMessageEventStream: a start event, incremental text events as +// deltas arrive, and a terminal done event carrying the aggregated message. +func (d *responsesDriver) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) { + if d.requiresAuth && strings.TrimSpace(req.Config.APIKey) == "" { + // Early "cannot build the stream": reference the provider, never a value. + return nil, fmt.Errorf("%s: missing API key", d.name) + } + + opts := make([]option.RequestOption, 0, len(d.clientOpts)+2) + if d.baseURL != "" { + opts = append(opts, option.WithBaseURL(d.baseURL)) + } + if key := strings.TrimSpace(req.Config.APIKey); key != "" { + opts = append(opts, option.WithAPIKey(key)) + } + opts = append(opts, d.clientOpts...) + client := openai.NewClient(opts...) + + params := buildResponsesParams(req) + + stream := NewAssistantMessageEventStream(0) + go d.pump(ctx, stream, &client, params) + return stream, nil +} + +// pump consumes the Responses SSE stream and translates events into pigo stream +// events. It always closes the stream. Every runtime failure (transport error, +// context cancellation, or an upstream error/failed event) becomes a terminal +// StreamErrorEvent (dual failure model), not a returned error. +// +// Incremental text.delta events emit a StreamTextEvent carrying the accumulated +// text so far, so the TUI renders tokens as they arrive. The terminal message is +// built from the authoritative response.completed payload via mapResponse, so +// the final aggregation matches the non-streamed result exactly. If no completed +// event arrives (a truncated stream that still ended cleanly), the accumulated +// delta text is used as a fallback. +func (d *responsesDriver) pump(ctx context.Context, stream *AssistantMessageEventStream, client *openai.Client, params responses.ResponseNewParams) { + defer stream.Close() + + if err := stream.Emit(ctx, StreamStartEvent{Partial: d.newPartial()}); err != nil { + return + } + + sse := client.Responses.NewStreaming(ctx, params) + defer sse.Close() + + var text strings.Builder + var thinking strings.Builder + var toolCalls []agentcore.ToolCallContent + var completed *responses.Response + for sse.Next() { + if ctx.Err() != nil { + d.emitError(stream, ctx.Err()) + return + } + switch variant := sse.Current().AsAny().(type) { + case responses.ResponseTextDeltaEvent: + text.WriteString(variant.Delta) + partial := d.buildPartial(thinking.String(), text.String(), toolCalls) + if err := stream.Emit(ctx, StreamTextEvent{Partial: partial}); err != nil { + return + } + case responses.ResponseReasoningSummaryTextDeltaEvent: + // The model's reasoning summary streams as its own text deltas, distinct + // from the answer text; accumulate it into a thinking block so the TUI + // renders reasoning the same way the chat driver does. + thinking.WriteString(variant.Delta) + partial := d.buildPartial(thinking.String(), text.String(), toolCalls) + if err := stream.Emit(ctx, StreamThinkingEvent{Partial: partial}); err != nil { + return + } + case responses.ResponseOutputItemDoneEvent: + // A finalized function_call item carries the model's tool request + // (name + arguments + call_id). Accumulate it and surface a tool-call + // partial so the TUI can show the pending call before the run ends. + if fc := variant.Item.AsFunctionCall(); fc.Type == "function_call" { + toolCalls = append(toolCalls, toolCallContent(fc)) + partial := d.buildPartial(thinking.String(), text.String(), toolCalls) + if err := stream.Emit(ctx, StreamToolCallEvent{Partial: partial}); err != nil { + return + } + } + case responses.ResponseCompletedEvent: + r := variant.Response + completed = &r + case responses.ResponseFailedEvent: + d.emitError(stream, fmt.Errorf("response failed")) + return + case responses.ResponseErrorEvent: + d.emitError(stream, fmt.Errorf("%s", variant.Message)) + return + } + } + if err := sse.Err(); err != nil { + d.emitError(stream, err) + return + } + + var msg agentcore.AssistantMessage + if completed != nil { + msg = d.mapResponse(completed) + } else { + msg = d.buildPartial(thinking.String(), text.String(), toolCalls) + msg.StopReason = agentcore.StopReasonEndTurn + if len(toolCalls) > 0 { + msg.StopReason = agentcore.StopReasonToolUse + } + } + stream.Emit(ctx, StreamDoneEvent{Message: msg}) +} + +// buildPartial assembles a cumulative snapshot message for a streaming partial: +// an optional thinking block (reasoning summary so far), the accumulated answer +// text, then any finalized tool calls — in the order the TUI should render them. +// All four emit sites in pump build partials through this one helper so they +// can't diverge. +func (d *responsesDriver) buildPartial(thinking, text string, toolCalls []agentcore.ToolCallContent) agentcore.AssistantMessage { + msg := d.newPartial() + if thinking != "" { + msg.Content = append(msg.Content, agentcore.NewThinkingContent(thinking)) + } + if text != "" { + msg.Content = append(msg.Content, agentcore.NewTextContent(text)) + } + msg.Content = appendToolCalls(msg.Content, toolCalls) + return msg +} + +// emitError emits a terminal StreamErrorEvent tagged for this provider. Uses a +// background context so the emit isn't dropped when ctx is already cancelled. +func (d *responsesDriver) emitError(stream *AssistantMessageEventStream, err error) { + stream.Emit(context.Background(), StreamErrorEvent{ + Message: agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + API: "openai", + Provider: d.name, + StopReason: agentcore.StopReasonError, + ErrorMessage: err.Error(), + }, + Err: fmt.Errorf("%s: %w", d.name, err), + }) +} + +// newPartial builds an empty assistant message tagged for this provider, the +// seed for start/text partials (mirrors OpenAIDecoder.partial()'s identity). +func (d *responsesDriver) newPartial() agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + API: "openai", + Provider: d.name, + } +} + +// mapResponse materializes a completed Responses API result into pigo's +// AssistantMessage: a reasoning summary (as a thinking block, when present), +// text content, tool calls, usage (when present), diagnostics, and a stop reason +// (tool_use when the model requested a tool, otherwise end_turn). +func (d *responsesDriver) mapResponse(resp *responses.Response) agentcore.AssistantMessage { + msg := d.newPartial() + msg.StopReason = agentcore.StopReasonEndTurn + msg.ResponseID = resp.ID + msg.ResponseModel = string(resp.Model) + if thinking := reasoningText(resp); thinking != "" { + msg.Content = append(msg.Content, agentcore.NewThinkingContent(thinking)) + } + if text := resp.OutputText(); text != "" { + msg.Content = append(msg.Content, agentcore.NewTextContent(text)) + } + var sawToolCall bool + for _, item := range resp.Output { + if fc := item.AsFunctionCall(); fc.Type == "function_call" { + msg.Content = append(msg.Content, toolCallContent(fc)) + sawToolCall = true + } + } + if sawToolCall { + msg.StopReason = agentcore.StopReasonToolUse + } + if resp.Usage.InputTokens != 0 || resp.Usage.OutputTokens != 0 { + msg.Usage = &agentcore.Usage{ + InputTokens: int(resp.Usage.InputTokens), + OutputTokens: int(resp.Usage.OutputTokens), + } + } + return msg +} + +// toolCallContent maps a Responses function_call item into a pigo +// ToolCallContent, keyed by the model's call_id so the tool result can be +// backfilled against it on the next turn. Arguments ride verbatim as raw JSON. +func toolCallContent(fc responses.ResponseFunctionToolCall) agentcore.ToolCallContent { + return agentcore.NewToolCallContent(fc.CallID, fc.Name, json.RawMessage(fc.Arguments)) +} + +// reasoningText concatenates the summary text of every reasoning item in a +// completed response. The Responses API returns the model's reasoning as one or +// more reasoning items, each carrying summary parts; pigo surfaces the joined +// text as a single thinking block, mirroring how the chat driver renders +// accumulated reasoning_content. +func reasoningText(resp *responses.Response) string { + var b strings.Builder + for _, item := range resp.Output { + if r := item.AsReasoning(); r.Type == "reasoning" { + for _, s := range r.Summary { + b.WriteString(s.Text) + } + } + } + return b.String() +} + +// appendToolCalls appends each accumulated tool call to a content list. Kept +// separate so the streaming partial and the terminal message build identical +// content from the same source. +func appendToolCalls(content agentcore.ContentList, calls []agentcore.ToolCallContent) agentcore.ContentList { + for _, c := range calls { + content = append(content, c) + } + return content +} + +// buildResponsesParams maps a CompletionRequest onto Responses API params. The +// system prompt becomes Instructions; the thinking level becomes a reasoning +// effort (with an auto summary so reasoning is returned); pigo tools become +// Responses function tools; and each message is replayed as the matching input +// item(s): assistant tool calls as function_call items, tool results as +// function_call_output items, and text (plus any images) as a role-tagged +// message. +func buildResponsesParams(req CompletionRequest) responses.ResponseNewParams { + params := responses.ResponseNewParams{ + Model: shared.ResponsesModel(req.Model), + } + if sp := strings.TrimSpace(req.Context.SystemPrompt); sp != "" { + params.Instructions = openai.String(sp) + } + if effort := responsesReasoningEffort(req.Config.ThinkingLevel); effort != "" { + // Requesting a summary makes the API return the model's reasoning so pigo + // can render it as a thinking block, matching the chat driver. + params.Reasoning = shared.ReasoningParam{Effort: effort, Summary: shared.ReasoningSummaryAuto} + } + if tools := buildResponsesTools(req.Context.Tools); len(tools) > 0 { + params.Tools = tools + } + + items := make(responses.ResponseInputParam, 0, len(req.Context.Messages)) + for _, m := range req.Context.Messages { + items = appendInputItems(items, m) + } + params.Input = responses.ResponseNewParamsInputUnion{OfInputItemList: items} + return params +} + +// buildResponsesTools converts pigo tools into Responses function tools. Each +// tool's JSON Schema becomes the function parameters; a schema that is empty or +// not a JSON object falls back to an empty object schema so the wire stays +// valid. Strict mode is off: pigo schemas are not authored against the Responses +// strict-function contract (which requires additionalProperties:false etc.). +func buildResponsesTools(tools []agentcore.AgentTool) []responses.ToolUnionParam { + if len(tools) == 0 { + return nil + } + out := make([]responses.ToolUnionParam, 0, len(tools)) + for _, t := range tools { + params := map[string]any{} + if raw := t.Schema(); len(raw) > 0 { + if err := json.Unmarshal(raw, ¶ms); err != nil { + params = map[string]any{} + } + } + tool := responses.ToolParamOfFunction(t.Name(), params, false) + if desc := t.Description(); desc != "" { + tool.OfFunction.Description = openai.String(desc) + } + out = append(out, tool) + } + return out +} + +// appendInputItems replays one pigo message as its Responses input item(s). +func appendInputItems(items responses.ResponseInputParam, m agentcore.Message) responses.ResponseInputParam { + switch msg := m.(type) { + case agentcore.ToolResultMessage: + // A tool result is backfilled against the model's call_id so the model + // can pair it with the request it issued the previous turn. + items = append(items, responses.ResponseInputItemParamOfFunctionCallOutput( + msg.ToolCallID, contentText(msg.Content))) + case agentcore.AssistantMessage: + if text := contentText(msg.Content); text != "" { + items = append(items, responses.ResponseInputItemParamOfMessage( + text, responses.EasyInputMessageRoleAssistant)) + } + for _, call := range msg.ToolCalls() { + items = append(items, responses.ResponseInputItemParamOfFunctionCall( + string(call.Arguments), call.ID, call.Name)) + } + default: + // A user (or other non-assistant) message with images is replayed as a + // content-part list (input_text + input_image data URIs); a text-only + // message stays a plain string. + if parts, ok := imageInputParts(m); ok { + items = append(items, responses.ResponseInputItemParamOfMessage(parts, responsesRole(m.Role()))) + } else if text := messageText(m); text != "" { + items = append(items, responses.ResponseInputItemParamOfMessage( + text, responsesRole(m.Role()))) + } + } + return items +} + +// imageInputParts builds a Responses content-part list for a message that +// carries at least one image: leading input_text (the concatenated text, if +// any) followed by one input_image per image, each as a data URI. It returns +// ok=false when the message has no images, so the caller keeps the plain-text +// path. +func imageInputParts(m agentcore.Message) (responses.ResponseInputMessageContentListParam, bool) { + content := messageContent(m) + var hasImage bool + for _, c := range content { + if _, ok := c.(agentcore.ImageContent); ok { + hasImage = true + break + } + } + if !hasImage { + return nil, false + } + parts := make(responses.ResponseInputMessageContentListParam, 0, len(content)+1) + if text := contentText(content); text != "" { + parts = append(parts, responses.ResponseInputContentParamOfInputText(text)) + } + for _, c := range content { + img, ok := c.(agentcore.ImageContent) + if !ok { + continue + } + part := responses.ResponseInputContentParamOfInputImage(responses.ResponseInputImageDetailAuto) + part.OfInputImage.ImageURL = openai.String(fmt.Sprintf("data:%s;base64,%s", img.MimeType, img.Data)) + parts = append(parts, part) + } + return parts, true +} + +// responsesReasoningEffort maps pigo's thinking level to a Responses API +// reasoning effort. The Responses reasoning field supports only low/medium/high, +// so "minimal" collapses to "low" (unlike the chat driver, which forwards +// "minimal" verbatim). off/unset yields "", signalling no reasoning param. +func responsesReasoningEffort(level agentcore.ThinkingLevel) shared.ReasoningEffort { + switch level { + case agentcore.ThinkingMinimal, agentcore.ThinkingLow: + return shared.ReasoningEffortLow + case agentcore.ThinkingMedium: + return shared.ReasoningEffortMedium + case agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax: + return shared.ReasoningEffortHigh + default: + return "" + } +} + +// responsesRole maps a pigo message role to the Responses API input role. Tool +// results are surfaced as user turns for this text milestone. +func responsesRole(role string) responses.EasyInputMessageRole { + switch role { + case agentcore.RoleAssistant: + return responses.EasyInputMessageRoleAssistant + default: + return responses.EasyInputMessageRoleUser + } +} + +// messageContent returns the content list of a message regardless of its +// concrete role type, so callers can inspect it for images. +func messageContent(m agentcore.Message) agentcore.ContentList { + switch msg := m.(type) { + case agentcore.UserMessage: + return msg.Content + case agentcore.AssistantMessage: + return msg.Content + case agentcore.ToolResultMessage: + return msg.Content + } + return nil +} + +// messageText concatenates the text blocks of a message, ignoring non-text +// content (handled in later milestones). +func messageText(m agentcore.Message) string { + var b strings.Builder + collectText(&b, messageContent(m)) + return b.String() +} + +func collectText(b *strings.Builder, content agentcore.ContentList) { + for _, c := range content { + if tc, ok := c.(agentcore.TextContent); ok { + b.WriteString(tc.Text) + } + } +} + +// contentText concatenates the text blocks of a content list. +func contentText(content agentcore.ContentList) string { + var b strings.Builder + collectText(&b, content) + return b.String() +} diff --git a/pigo/internal/provider/responses_test.go b/pigo/internal/provider/responses_test.go new file mode 100644 index 0000000..b861a0d --- /dev/null +++ b/pigo/internal/provider/responses_test.go @@ -0,0 +1,756 @@ +package provider + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "strconv" + "strings" + "testing" + + "github.com/openai/openai-go/option" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// roundTripFunc adapts a function to http.RoundTripper so a test can stub the +// SDK transport without a live endpoint. +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +// newResponsesTestDriver builds a resp_api driver whose SDK client is pointed at +// the given stub round-tripper, capturing the request path the SDK targets. +func newResponsesTestDriver(baseURL string, rt roundTripFunc) *responsesDriver { + d := NewOpenAIResponsesProvider("openai", baseURL, nil) + d.clientOpts = []option.RequestOption{ + option.WithHTTPClient(&http.Client{Transport: rt}), + } + return d +} + +func jsonResponse(status int, body string) *http.Response { + return &http.Response{ + StatusCode: status, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +// sseResponse builds a 200 text/event-stream response whose body is the given +// SSE data frames, mirroring how the Responses API streams events. Each frame is +// a JSON object carrying its own "type" discriminator. +func sseResponse(frames ...string) *http.Response { + var b strings.Builder + for _, f := range frames { + b.WriteString("data: ") + b.WriteString(f) + b.WriteString("\n\n") + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(b.String())), + } +} + +// completedFrame is a response.completed SSE frame whose embedded Response yields +// the given output text, id, model, and token usage — the authoritative terminal +// payload the driver maps into its final message. +func completedFrame(text, id, model string, inTok, outTok int) string { + return `{"type":"response.completed","sequence_number":99,"response":{` + + `"id":"` + id + `","model":"` + model + `",` + + `"output":[{"type":"message","role":"assistant","status":"completed",` + + `"content":[{"type":"output_text","text":"` + text + `"}]}],` + + `"usage":{"input_tokens":` + itoa(inTok) + `,"output_tokens":` + itoa(outTok) + + `,"input_tokens_details":{"cached_tokens":0},"output_tokens_details":{"reasoning_tokens":0}}}}` +} + +func deltaFrame(delta string) string { + return `{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,` + + `"content_index":0,"sequence_number":1,"logprobs":[],"delta":"` + delta + `"}` +} + +func itoa(n int) string { return strconv.Itoa(n) } + +// drain collects the terminal message from a stream, mirroring how the loop +// consumes a provider stream. +func drain(t *testing.T, stream *AssistantMessageEventStream) agentcore.AssistantMessage { + t.Helper() + for range stream.Events() { + } + msg, err := stream.Result(context.Background()) + if err != nil { + t.Fatalf("stream result error: %v", err) + } + return msg +} + +func userMsg(text string) agentcore.UserMessage { + return agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + } +} + +func TestResponsesDriverPostsToResponsesEndpoint(t *testing.T) { + var gotPath, gotBody string + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + gotPath = r.URL.Path + if r.Body != nil { + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + } + return sseResponse( + deltaFrame("hi "), + deltaFrame("there"), + completedFrame("hi there", "resp_123", "gpt-4o", 11, 7), + ), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + req := CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{ + SystemPrompt: "be terse", + Messages: agentcore.MessageList{userMsg("hello")}, + }, + Config: StreamConfig{APIKey: "sk-test"}, + } + stream, err := d.StreamCompletion(context.Background(), req) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + msg := drain(t, stream) + + if !strings.HasSuffix(gotPath, "/responses") { + t.Errorf("request path = %q, want to end with /responses", gotPath) + } + // The prompt and system instruction must reach the wire body. + if !strings.Contains(gotBody, "hello") { + t.Errorf("request body missing prompt: %q", gotBody) + } + var payload map[string]any + if err := json.Unmarshal([]byte(gotBody), &payload); err != nil { + t.Fatalf("request body not valid JSON: %v", err) + } + if payload["instructions"] != "be terse" { + t.Errorf("instructions = %v, want %q", payload["instructions"], "be terse") + } + if payload["model"] != "gpt-4o" { + t.Errorf("model = %v, want gpt-4o", payload["model"]) + } + // A streaming call must set stream:true on the wire. + if payload["stream"] != true { + t.Errorf("stream = %v, want true", payload["stream"]) + } + + if got := textOf(msg); got != "hi there" { + t.Errorf("assistant text = %q, want %q", got, "hi there") + } + if msg.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("stop reason = %q, want end_turn", msg.StopReason) + } + if msg.ResponseID != "resp_123" { + t.Errorf("response id = %q, want resp_123", msg.ResponseID) + } + if msg.Usage == nil || msg.Usage.InputTokens != 11 || msg.Usage.OutputTokens != 7 { + t.Errorf("usage = %+v, want {11 7}", msg.Usage) + } + if msg.API != "openai" || msg.Provider != "openai" { + t.Errorf("tags = api:%q provider:%q, want openai/openai", msg.API, msg.Provider) + } +} + +// The driver must emit incremental text partials as deltas arrive, and each +// partial must carry the text accumulated so far (not just the latest delta), so +// the terminal message equals the concatenation the caller already rendered. +func TestResponsesDriverStreamsIncrementalDeltas(t *testing.T) { + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + return sseResponse( + deltaFrame("Hello"), + deltaFrame(", "), + deltaFrame("world"), + completedFrame("Hello, world", "resp_9", "gpt-4o", 3, 4), + ), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{userMsg("hi")}}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + + var textPartials []string + for ev := range stream.Events() { + if te, ok := ev.(StreamTextEvent); ok { + textPartials = append(textPartials, textOf(te.Partial)) + } + } + msg, err := stream.Result(context.Background()) + if err != nil { + t.Fatalf("stream result error: %v", err) + } + + want := []string{"Hello", "Hello, ", "Hello, world"} + if len(textPartials) != len(want) { + t.Fatalf("got %d text partials %q, want %d %q", len(textPartials), textPartials, len(want), want) + } + for i := range want { + if textPartials[i] != want[i] { + t.Errorf("partial[%d] = %q, want %q", i, textPartials[i], want[i]) + } + } + // Final aggregation must match the completed payload, i.e. the last partial. + if got := textOf(msg); got != "Hello, world" { + t.Errorf("final text = %q, want %q", got, "Hello, world") + } +} + +// A cancelled context must terminate the stream with an error rather than +// yielding a normal end_turn message. The transport cancels mid-flight (after +// the stream has started) and reports the cancellation, mirroring how an +// in-progress SSE read aborts when the caller cancels. +func TestResponsesDriverContextCancelStopsStream(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + cancel() + return nil, context.Canceled + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + stream, err := d.StreamCompletion(ctx, CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{userMsg("hi")}}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion should not early-error on cancel: %v", err) + } + + var sawError bool + for ev := range stream.Events() { + if _, ok := ev.(StreamErrorEvent); ok { + sawError = true + } + } + if !sawError { + t.Fatal("expected a terminal StreamErrorEvent after context cancel") + } + msg, _ := stream.Result(context.Background()) + if msg.StopReason != agentcore.StopReasonError { + t.Errorf("stop reason = %q, want error", msg.StopReason) + } +} + +// A non-2xx from the endpoint must ride the stream as a terminal error event, +// not be returned from StreamCompletion (dual failure model, FR-13). +func TestResponsesDriverUpstreamErrorRidesStream(t *testing.T) { + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + return jsonResponse(http.StatusUnauthorized, `{"error":{"message":"bad key"}}`), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + req := CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{userMsg("hello")}}, + Config: StreamConfig{APIKey: "sk-test"}, + } + stream, err := d.StreamCompletion(context.Background(), req) + if err != nil { + t.Fatalf("StreamCompletion should not early-error on upstream failure: %v", err) + } + + var sawError bool + for ev := range stream.Events() { + if _, ok := ev.(StreamErrorEvent); ok { + sawError = true + } + } + if !sawError { + t.Fatal("expected a terminal StreamErrorEvent for a 401 response") + } + msg, _ := stream.Result(context.Background()) + if msg.StopReason != agentcore.StopReasonError { + t.Errorf("stop reason = %q, want error", msg.StopReason) + } +} + +// An in-band error event (type "error") mid-stream must ride the stream as a +// terminal error, carrying the event's message. This is a distinct path from a +// transport-level non-2xx (which surfaces via the stream's Err()). +func TestResponsesDriverInStreamErrorEvent(t *testing.T) { + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + return sseResponse( + deltaFrame("partial"), + `{"type":"error","code":"server_error","message":"boom","param":"","sequence_number":2}`, + ), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{userMsg("hi")}}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion should not early-error: %v", err) + } + + var errEvent *StreamErrorEvent + for ev := range stream.Events() { + if se, ok := ev.(StreamErrorEvent); ok { + e := se + errEvent = &e + } + } + if errEvent == nil { + t.Fatal("expected a terminal StreamErrorEvent for an in-band error event") + } + if !strings.Contains(errEvent.Message.ErrorMessage, "boom") { + t.Errorf("error message = %q, want to contain %q", errEvent.Message.ErrorMessage, "boom") + } + msg, _ := stream.Result(context.Background()) + if msg.StopReason != agentcore.StopReasonError { + t.Errorf("stop reason = %q, want error", msg.StopReason) + } +} + +// A missing API key is the one early "cannot build the stream" error. +func TestResponsesDriverMissingKeyIsEarlyError(t *testing.T) { + d := NewOpenAIResponsesProvider("openai", "https://api.openai.test/v1", nil) + _, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Config: StreamConfig{APIKey: " "}, + }) + if err == nil { + t.Fatal("expected early error for missing API key") + } + if !strings.Contains(err.Error(), "missing API key") { + t.Errorf("error = %q, want to mention missing API key", err.Error()) + } +} + +// textOf returns the concatenated text content of an assistant message. +func textOf(m agentcore.AssistantMessage) string { + var b bytes.Buffer + for _, c := range m.Content { + if tc, ok := c.(agentcore.TextContent); ok { + b.WriteString(tc.Text) + } + } + return b.String() +} + +// fakeTool is a minimal AgentTool for exercising tool-schema serialization; it +// never executes in these transport-level tests. +type fakeTool struct { + name string + desc string + schema json.RawMessage +} + +func (t fakeTool) Name() string { return t.name } +func (t fakeTool) Description() string { return t.desc } +func (t fakeTool) Schema() json.RawMessage { return t.schema } +func (t fakeTool) ExecutionMode() agentcore.ToolExecutionMode { return agentcore.ToolExecutionParallel } +func (t fakeTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{}, nil +} + +// functionCallDoneFrame is a response.output_item.done SSE frame carrying a +// finalized function_call item (the model's tool request). +func functionCallDoneFrame(callID, name, argsJSON string) string { + frame := map[string]any{ + "type": "response.output_item.done", + "output_index": 0, + "sequence_number": 5, + "item": map[string]any{ + "type": "function_call", + "id": "fc_1", + "call_id": callID, + "name": name, + "arguments": argsJSON, + "status": "completed", + }, + } + b, _ := json.Marshal(frame) + return string(b) +} + +// completedToolFrame is a response.completed frame whose output is a single +// function_call item (no assistant text) plus token usage. +func completedToolFrame(id, model, callID, name, argsJSON string) string { + frame := map[string]any{ + "type": "response.completed", + "sequence_number": 99, + "response": map[string]any{ + "id": id, + "model": model, + "output": []any{map[string]any{ + "type": "function_call", + "id": "fc_1", + "call_id": callID, + "name": name, + "arguments": argsJSON, + "status": "completed", + }}, + "usage": map[string]any{ + "input_tokens": 5, + "output_tokens": 2, + "input_tokens_details": map[string]any{"cached_tokens": 0}, + "output_tokens_details": map[string]any{"reasoning_tokens": 0}, + }, + }, + } + b, _ := json.Marshal(frame) + return string(b) +} + +// A pigo tool must reach the wire as a Responses function tool: type "function", +// its name, JSON-Schema parameters, and description. +func TestResponsesDriverSendsToolSchema(t *testing.T) { + var gotBody string + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + if r.Body != nil { + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + } + return sseResponse(completedFrame("ok", "resp_1", "gpt-4o", 1, 1)), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + tool := fakeTool{ + name: "read_file", + desc: "reads a file", + schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}`), + } + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{ + Messages: agentcore.MessageList{userMsg("read a.go")}, + Tools: []agentcore.AgentTool{tool}, + }, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + drain(t, stream) + + var payload map[string]any + if err := json.Unmarshal([]byte(gotBody), &payload); err != nil { + t.Fatalf("request body not valid JSON: %v", err) + } + tools, ok := payload["tools"].([]any) + if !ok || len(tools) != 1 { + t.Fatalf("tools = %v, want a single-element array", payload["tools"]) + } + tool0 := tools[0].(map[string]any) + if tool0["type"] != "function" { + t.Errorf("tool type = %v, want function", tool0["type"]) + } + if tool0["name"] != "read_file" { + t.Errorf("tool name = %v, want read_file", tool0["name"]) + } + if tool0["description"] != "reads a file" { + t.Errorf("tool description = %v, want %q", tool0["description"], "reads a file") + } + params, ok := tool0["parameters"].(map[string]any) + if !ok { + t.Fatalf("tool parameters missing or not an object: %v", tool0["parameters"]) + } + props, ok := params["properties"].(map[string]any) + if !ok || props["path"] == nil { + t.Errorf("tool parameters.properties.path missing: %v", params) + } +} + +// A function_call in the completed response must be parsed into a pigo +// ToolCallContent (id + name + raw arguments) and set StopReason=tool_use; a +// StreamToolCallEvent must also surface the pending call mid-stream. +func TestResponsesDriverParsesToolCall(t *testing.T) { + args := `{"path":"a.go"}` + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + return sseResponse( + functionCallDoneFrame("call_abc", "read_file", args), + completedToolFrame("resp_7", "gpt-4o", "call_abc", "read_file", args), + ), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{userMsg("read a.go")}}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + + var sawToolCallEvent bool + for ev := range stream.Events() { + if _, ok := ev.(StreamToolCallEvent); ok { + sawToolCallEvent = true + } + } + if !sawToolCallEvent { + t.Error("expected a StreamToolCallEvent mid-stream") + } + msg, err := stream.Result(context.Background()) + if err != nil { + t.Fatalf("stream result error: %v", err) + } + + calls := msg.ToolCalls() + if len(calls) != 1 { + t.Fatalf("got %d tool calls, want 1", len(calls)) + } + if calls[0].ID != "call_abc" || calls[0].Name != "read_file" { + t.Errorf("tool call = id:%q name:%q, want call_abc/read_file", calls[0].ID, calls[0].Name) + } + if string(calls[0].Arguments) != args { + t.Errorf("tool call arguments = %q, want %q", calls[0].Arguments, args) + } + if msg.StopReason != agentcore.StopReasonToolUse { + t.Errorf("stop reason = %q, want tool_use", msg.StopReason) + } +} + +// On a follow-up turn, a prior assistant tool call must be replayed as a +// function_call input item and its result as a function_call_output item, both +// keyed by the same call_id, so the model can pair request and result. +func TestResponsesDriverBackfillsToolResult(t *testing.T) { + var gotBody string + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + if r.Body != nil { + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + } + return sseResponse(completedFrame("done", "resp_2", "gpt-4o", 8, 3)), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + assistant := agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{ + agentcore.NewToolCallContent("call_abc", "read_file", json.RawMessage(`{"path":"a.go"}`)), + }, + } + result := agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, + ToolCallID: "call_abc", + ToolName: "read_file", + Content: agentcore.ContentList{agentcore.NewTextContent("package main")}, + } + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{ + userMsg("read a.go"), assistant, result, + }}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + drain(t, stream) + + var payload struct { + Input []map[string]any `json:"input"` + } + if err := json.Unmarshal([]byte(gotBody), &payload); err != nil { + t.Fatalf("request body not valid JSON: %v", err) + } + var sawCall, sawOutput bool + for _, item := range payload.Input { + switch item["type"] { + case "function_call": + sawCall = true + if item["call_id"] != "call_abc" || item["name"] != "read_file" { + t.Errorf("function_call item = %v, want call_abc/read_file", item) + } + if item["arguments"] != `{"path":"a.go"}` { + t.Errorf("function_call arguments = %v, want the raw args JSON string", item["arguments"]) + } + case "function_call_output": + sawOutput = true + if item["call_id"] != "call_abc" { + t.Errorf("function_call_output call_id = %v, want call_abc", item["call_id"]) + } + if item["output"] != "package main" { + t.Errorf("function_call_output output = %v, want %q", item["output"], "package main") + } + } + } + if !sawCall { + t.Error("wire input missing the replayed function_call item") + } + if !sawOutput { + t.Error("wire input missing the function_call_output item") + } +} + +// completedReasoningFrame is a response.completed frame whose output carries a +// reasoning item (summary text) followed by the assistant message text. +func completedReasoningFrame(id, model, summary, text string) string { + frame := map[string]any{ + "type": "response.completed", + "sequence_number": 99, + "response": map[string]any{ + "id": id, + "model": model, + "output": []any{ + map[string]any{ + "type": "reasoning", + "id": "rs_1", + "summary": []any{map[string]any{ + "type": "summary_text", + "text": summary, + }}, + }, + map[string]any{ + "type": "message", + "role": "assistant", + "status": "completed", + "content": []any{map[string]any{ + "type": "output_text", + "text": text, + }}, + }, + }, + "usage": map[string]any{ + "input_tokens": 3, + "output_tokens": 4, + "input_tokens_details": map[string]any{"cached_tokens": 0}, + "output_tokens_details": map[string]any{"reasoning_tokens": 2}, + }, + }, + } + b, _ := json.Marshal(frame) + return string(b) +} + +// A user message carrying an image must reach the wire as a message whose +// content is a part list: an input_text part plus an input_image part whose +// image_url is the base64 data URI. +func TestResponsesDriverSendsImageInput(t *testing.T) { + var gotBody string + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + if r.Body != nil { + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + } + return sseResponse(completedFrame("ok", "resp_1", "gpt-4o", 1, 1)), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + imgMsg := agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{ + agentcore.NewTextContent("what is this?"), + agentcore.NewImageContent("aGVsbG8=", "image/png"), + }, + } + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{imgMsg}}, + Config: StreamConfig{APIKey: "sk-test"}, + }) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + drain(t, stream) + + var payload struct { + Input []struct { + Type string `json:"type"` + Role string `json:"role"` + Content []map[string]any `json:"content"` + } `json:"input"` + } + if err := json.Unmarshal([]byte(gotBody), &payload); err != nil { + t.Fatalf("request body not valid JSON: %v", err) + } + if len(payload.Input) != 1 { + t.Fatalf("got %d input items, want 1", len(payload.Input)) + } + parts := payload.Input[0].Content + var sawText, sawImage bool + for _, p := range parts { + switch p["type"] { + case "input_text": + sawText = true + if p["text"] != "what is this?" { + t.Errorf("input_text = %v, want %q", p["text"], "what is this?") + } + case "input_image": + sawImage = true + if p["image_url"] != "data:image/png;base64,aGVsbG8=" { + t.Errorf("input_image image_url = %v, want the data URI", p["image_url"]) + } + } + } + if !sawText { + t.Error("wire input missing the input_text part") + } + if !sawImage { + t.Error("wire input missing the input_image part") + } +} + +// A request with a thinking level must set the reasoning.effort (and an auto +// summary) on the wire, and a reasoning item in the completed response must be +// parsed into a leading ThinkingContent block. +func TestResponsesDriverReasoning(t *testing.T) { + var gotBody string + rt := roundTripFunc(func(r *http.Request) (*http.Response, error) { + if r.Body != nil { + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + } + return sseResponse(completedReasoningFrame("resp_9", "gpt-4o", "let me think", "the answer")), nil + }) + d := newResponsesTestDriver("https://api.openai.test/v1", rt) + + stream, err := d.StreamCompletion(context.Background(), CompletionRequest{ + Model: "gpt-4o", + Context: LlmContext{Messages: agentcore.MessageList{userMsg("solve it")}}, + Config: StreamConfig{APIKey: "sk-test", ThinkingLevel: agentcore.ThinkingMedium}, + }) + if err != nil { + t.Fatalf("StreamCompletion returned early error: %v", err) + } + msg := drain(t, stream) + + var payload struct { + Reasoning struct { + Effort string `json:"effort"` + Summary string `json:"summary"` + } `json:"reasoning"` + } + if err := json.Unmarshal([]byte(gotBody), &payload); err != nil { + t.Fatalf("request body not valid JSON: %v", err) + } + if payload.Reasoning.Effort != "medium" { + t.Errorf("reasoning.effort = %q, want medium", payload.Reasoning.Effort) + } + if payload.Reasoning.Summary != "auto" { + t.Errorf("reasoning.summary = %q, want auto", payload.Reasoning.Summary) + } + + var thinking string + for _, c := range msg.Content { + if tc, ok := c.(agentcore.ThinkingContent); ok { + thinking = tc.Thinking + } + } + if thinking != "let me think" { + t.Errorf("thinking content = %q, want %q", thinking, "let me think") + } +} diff --git a/pigo/internal/provider/special_auth.go b/pigo/internal/provider/special_auth.go new file mode 100644 index 0000000..f259e1c --- /dev/null +++ b/pigo/internal/provider/special_auth.go @@ -0,0 +1,225 @@ +// This file implements parameter validation and endpoint construction for the +// "special auth" providers (US-007 / FR-12): Azure OpenAI, Amazon Bedrock, +// Google Vertex, and Cloudflare (Workers AI + AI Gateway). Unlike the generic +// bearer/x-api-key providers, each of these composes its endpoint from several +// environment variables and/or needs a non-standard credential path, so a +// dedicated resolver validates the required parameters and builds the concrete +// base URL before handing off to the shared OpenAI-/Anthropic-compatible driver. +// +// Scope note (PRD Non-Goals): AWS SigV4 request signing is NOT implemented. +// Bedrock supports only the AWS_BEARER_TOKEN_BEDROCK bearer path; other AWS +// credential sources (AWS_PROFILE, AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY) +// are only *detected* so that a clear, actionable error is returned instead of +// an opaque auth failure. +// +// Security: this file reads env var NAMES and composes URLs from non-secret +// parameters (region, resource name, account id, …). Secret values (API keys, +// bearer tokens) are never logged or embedded in error text — errors name the +// absent env var, never a value. +package provider + +import ( + "fmt" + "strings" +) + +// IsSpecialAuthProvider reports whether a provider spec needs the bespoke +// endpoint-construction / credential-validation handled by ResolveSpecialProvider, +// rather than the generic driver wiring. It matches the multi-parameter auth +// schemes (azure/aws/special) and the two Cloudflare providers (which keep a +// standard auth scheme but still compose their endpoint from env vars). +func IsSpecialAuthProvider(spec ProviderSpec) bool { + switch spec.AuthScheme { + case AuthAzure, AuthAWS, AuthSpecial: + return true + } + return strings.HasPrefix(spec.Name, "cloudflare-") +} + +// ResolveSpecialProvider validates the required parameters for a special-auth +// provider and constructs the matching wire driver against the composed base +// URL. flagBaseURL is the explicit --base-url override (highest precedence, wins +// over any composed default); env resolves environment variables (os.Getenv in +// production, a fake in tests). A missing required parameter yields an error +// naming exactly which env var is absent; no network request is made here. +func ResolveSpecialProvider(spec ProviderSpec, model, flagBaseURL string, env func(string) string) (Provider, error) { + if env == nil { + env = func(string) string { return "" } + } + models := []Model{{Provider: spec.Name, ID: model, SupportsImages: true}} + switch spec.Name { + case "azure-openai-responses": + return resolveAzureOpenAI(spec, model, flagBaseURL, env, models) + case "amazon-bedrock": + return resolveBedrock(spec, flagBaseURL, env, models) + case "google-vertex": + return resolveGoogleVertex(spec, flagBaseURL, env, models) + case "cloudflare-workers-ai": + return resolveCloudflareWorkersAI(spec, flagBaseURL, env, models) + case "cloudflare-ai-gateway": + return resolveCloudflareAIGateway(spec, flagBaseURL, env, models) + default: + return nil, fmt.Errorf("provider %q is not a special-auth provider", spec.Name) + } +} + +// resolveAzureOpenAI composes the Azure OpenAI endpoint. The endpoint origin is +// AZURE_OPENAI_BASE_URL (or the --base-url override), else it is built from +// AZURE_OPENAI_RESOURCE_NAME as https://{resource}.openai.azure.com. The API +// version (AZURE_OPENAI_API_VERSION, default "v1") and an optional deployment +// mapping (AZURE_OPENAI_DEPLOYMENT_NAME_MAP) shape the path. Auth uses +// AZURE_OPENAI_API_KEY over the OpenAI wire. +func resolveAzureOpenAI(_ ProviderSpec, model, flagBaseURL string, env func(string) string, models []Model) (Provider, error) { + if strings.TrimSpace(env("AZURE_OPENAI_API_KEY")) == "" { + return nil, fmt.Errorf("azure-openai-responses: missing required env var AZURE_OPENAI_API_KEY") + } + origin := strings.TrimSpace(flagBaseURL) + if origin == "" { + origin = strings.TrimSpace(env("AZURE_OPENAI_BASE_URL")) + } + if origin == "" { + resource := strings.TrimSpace(env("AZURE_OPENAI_RESOURCE_NAME")) + if resource == "" { + return nil, fmt.Errorf("azure-openai-responses: missing endpoint configuration; set AZURE_OPENAI_BASE_URL or AZURE_OPENAI_RESOURCE_NAME") + } + origin = fmt.Sprintf("https://%s.openai.azure.com", resource) + } + apiVersion := strings.TrimSpace(env("AZURE_OPENAI_API_VERSION")) + if apiVersion == "" { + apiVersion = "v1" + } + deployment := resolveAzureDeployment(env("AZURE_OPENAI_DEPLOYMENT_NAME_MAP"), model) + baseURL := azureEndpoint(origin, apiVersion, deployment) + return NewOpenAICompatibleProvider(baseURL, models), nil +} + +// azureEndpoint builds the Azure OpenAI base URL from a validated origin. When a +// deployment is resolved for the model, the classic deployment-scoped path is +// used (…/openai/deployments/{deployment}); otherwise the version-scoped v1 path +// (…/openai/{apiVersion}) is used. The shared driver appends /chat/completions. +func azureEndpoint(origin, apiVersion, deployment string) string { + origin = strings.TrimRight(strings.TrimSpace(origin), "/") + if deployment != "" { + return fmt.Sprintf("%s/openai/deployments/%s", origin, deployment) + } + return fmt.Sprintf("%s/openai/%s", origin, apiVersion) +} + +// resolveAzureDeployment parses AZURE_OPENAI_DEPLOYMENT_NAME_MAP (a +// comma-separated list of model=deployment pairs) and returns the deployment +// mapped to model, or "" when the map is empty or has no entry for the model. +func resolveAzureDeployment(raw, model string) string { + m := parseDeploymentMap(raw) + return m[strings.TrimSpace(model)] +} + +// parseDeploymentMap parses a comma-separated "model=deployment" list into a +// map. Blank entries and entries without '=' are skipped; keys and values are +// trimmed. It never returns nil so lookups are always safe. +func parseDeploymentMap(raw string) map[string]string { + out := make(map[string]string) + for _, pair := range strings.Split(raw, ",") { + pair = strings.TrimSpace(pair) + if pair == "" { + continue + } + k, v, ok := strings.Cut(pair, "=") + k, v = strings.TrimSpace(k), strings.TrimSpace(v) + if !ok || k == "" || v == "" { + continue + } + out[k] = v + } + return out +} + +// resolveBedrock composes the Amazon Bedrock runtime endpoint +// (https://bedrock-runtime.{region}.amazonaws.com; region defaults to +// us-east-1) and validates credentials. Only the AWS_BEARER_TOKEN_BEDROCK +// bearer path is supported (SigV4 is out of scope): if only AWS_PROFILE or +// static AWS keys are present, a clear error explains SigV4 is unsupported and +// names the missing AWS_BEARER_TOKEN_BEDROCK. +func resolveBedrock(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) { + if strings.TrimSpace(env("AWS_BEARER_TOKEN_BEDROCK")) == "" { + hasProfile := strings.TrimSpace(env("AWS_PROFILE")) != "" + hasStaticKeys := strings.TrimSpace(env("AWS_ACCESS_KEY_ID")) != "" && + strings.TrimSpace(env("AWS_SECRET_ACCESS_KEY")) != "" + if hasProfile || hasStaticKeys { + return nil, fmt.Errorf("amazon-bedrock: detected AWS credentials (AWS_PROFILE / AWS_ACCESS_KEY_ID) but SigV4 request signing is not supported yet; set AWS_BEARER_TOKEN_BEDROCK to use the bearer-token path") + } + return nil, fmt.Errorf("amazon-bedrock: missing required env var AWS_BEARER_TOKEN_BEDROCK") + } + baseURL := strings.TrimSpace(flagBaseURL) + if baseURL == "" { + region := strings.TrimSpace(env("AWS_REGION")) + if region == "" { + region = "us-east-1" + } + baseURL = fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com", region) + } + return NewBedrockProvider(baseURL, models), nil +} + +// resolveGoogleVertex composes the Vertex AI endpoint +// (https://{location}-aiplatform.googleapis.com) and validates that a project, +// a location, and a credential source (GOOGLE_CLOUD_API_KEY or ADC via +// GOOGLE_APPLICATION_CREDENTIALS) are present, naming any absent env var. +func resolveGoogleVertex(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) { + if strings.TrimSpace(env("GOOGLE_CLOUD_PROJECT")) == "" { + return nil, fmt.Errorf("google-vertex: missing required env var GOOGLE_CLOUD_PROJECT") + } + location := strings.TrimSpace(env("GOOGLE_CLOUD_LOCATION")) + if location == "" { + return nil, fmt.Errorf("google-vertex: missing required env var GOOGLE_CLOUD_LOCATION") + } + if strings.TrimSpace(env("GOOGLE_CLOUD_API_KEY")) == "" && + strings.TrimSpace(env("GOOGLE_APPLICATION_CREDENTIALS")) == "" { + return nil, fmt.Errorf("google-vertex: missing credentials; set GOOGLE_CLOUD_API_KEY or GOOGLE_APPLICATION_CREDENTIALS (ADC)") + } + baseURL := strings.TrimSpace(flagBaseURL) + if baseURL == "" { + baseURL = fmt.Sprintf("https://%s-aiplatform.googleapis.com", location) + } + return NewOpenAICompatibleProvider(baseURL, models), nil +} + +// resolveCloudflareWorkersAI composes the Workers AI endpoint +// (https://api.cloudflare.com/client/v4/accounts/{account}/ai/v1), requiring +// CLOUDFLARE_API_KEY and CLOUDFLARE_ACCOUNT_ID. OpenAI wire. +func resolveCloudflareWorkersAI(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) { + if strings.TrimSpace(env("CLOUDFLARE_API_KEY")) == "" { + return nil, fmt.Errorf("cloudflare-workers-ai: missing required env var CLOUDFLARE_API_KEY") + } + account := strings.TrimSpace(env("CLOUDFLARE_ACCOUNT_ID")) + if account == "" { + return nil, fmt.Errorf("cloudflare-workers-ai: missing required env var CLOUDFLARE_ACCOUNT_ID") + } + baseURL := strings.TrimSpace(flagBaseURL) + if baseURL == "" { + baseURL = fmt.Sprintf("https://api.cloudflare.com/client/v4/accounts/%s/ai/v1", account) + } + return NewOpenAICompatibleProvider(baseURL, models), nil +} + +// resolveCloudflareAIGateway composes the AI Gateway endpoint +// (https://gateway.ai.cloudflare.com/v1/{account}/{gateway}/anthropic), +// requiring CLOUDFLARE_API_KEY, CLOUDFLARE_ACCOUNT_ID, and CLOUDFLARE_GATEWAY_ID. +// Anthropic wire. +func resolveCloudflareAIGateway(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) { + if strings.TrimSpace(env("CLOUDFLARE_API_KEY")) == "" { + return nil, fmt.Errorf("cloudflare-ai-gateway: missing required env var CLOUDFLARE_API_KEY") + } + account := strings.TrimSpace(env("CLOUDFLARE_ACCOUNT_ID")) + if account == "" { + return nil, fmt.Errorf("cloudflare-ai-gateway: missing required env var CLOUDFLARE_ACCOUNT_ID") + } + gateway := strings.TrimSpace(env("CLOUDFLARE_GATEWAY_ID")) + if gateway == "" { + return nil, fmt.Errorf("cloudflare-ai-gateway: missing required env var CLOUDFLARE_GATEWAY_ID") + } + baseURL := strings.TrimSpace(flagBaseURL) + if baseURL == "" { + baseURL = fmt.Sprintf("https://gateway.ai.cloudflare.com/v1/%s/%s/anthropic", account, gateway) + } + return NewAnthropicProvider(baseURL, models), nil +} diff --git a/pigo/internal/provider/special_auth_test.go b/pigo/internal/provider/special_auth_test.go new file mode 100644 index 0000000..62e82f0 --- /dev/null +++ b/pigo/internal/provider/special_auth_test.go @@ -0,0 +1,291 @@ +package provider + +import ( + "strings" + "testing" +) + +// envFrom builds an env(string)string lookup from a map for hermetic tests: no +// process environment is read, so tests never depend on ambient state and make +// no network requests. +func envFrom(m map[string]string) func(string) string { + return func(k string) string { return m[k] } +} + +// baseURLOf extracts the composed base URL from a constructed driver by type +// asserting the two concrete driver shapes (same-package access to unexported +// fields). It fails the test if the provider is neither shape. +func baseURLOf(t *testing.T, p Provider) string { + t.Helper() + switch d := p.(type) { + case *openAICompatDriver: + return d.baseURL + case *anthropicCompatDriver: + return d.baseURL + default: + t.Fatalf("unexpected provider type %T", p) + return "" + } +} + +func specFor(t *testing.T, name string) ProviderSpec { + t.Helper() + spec, ok := LookupProviderSpec(name) + if !ok { + t.Fatalf("registry missing provider %q", name) + } + return spec +} + +func TestResolveSpecialProvider_Azure(t *testing.T) { + spec := specFor(t, "azure-openai-responses") + + // Missing API key. + if _, err := ResolveSpecialProvider(spec, "gpt-4o", "", envFrom(nil)); err == nil || + !strings.Contains(err.Error(), "AZURE_OPENAI_API_KEY") { + t.Fatalf("expected AZURE_OPENAI_API_KEY error, got %v", err) + } + + // Key present but no endpoint origin. + env := envFrom(map[string]string{"AZURE_OPENAI_API_KEY": "k"}) + if _, err := ResolveSpecialProvider(spec, "gpt-4o", "", env); err == nil || + !strings.Contains(err.Error(), "AZURE_OPENAI_BASE_URL") || + !strings.Contains(err.Error(), "AZURE_OPENAI_RESOURCE_NAME") { + t.Fatalf("expected endpoint-config error naming both env vars, got %v", err) + } + + // Resource name → composed origin, default api version v1. + env = envFrom(map[string]string{ + "AZURE_OPENAI_API_KEY": "k", + "AZURE_OPENAI_RESOURCE_NAME": "myres", + }) + p, err := ResolveSpecialProvider(spec, "gpt-4o", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://myres.openai.azure.com/openai/v1"; got != want { + t.Fatalf("azure base_url = %q, want %q", got, want) + } + + // Explicit base URL env + custom api version. + env = envFrom(map[string]string{ + "AZURE_OPENAI_API_KEY": "k", + "AZURE_OPENAI_BASE_URL": "https://custom.example.com", + "AZURE_OPENAI_API_VERSION": "2024-10-01", + }) + p, err = ResolveSpecialProvider(spec, "gpt-4o", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://custom.example.com/openai/2024-10-01"; got != want { + t.Fatalf("azure base_url = %q, want %q", got, want) + } + + // Deployment name map → deployment-scoped path. + env = envFrom(map[string]string{ + "AZURE_OPENAI_API_KEY": "k", + "AZURE_OPENAI_RESOURCE_NAME": "myres", + "AZURE_OPENAI_DEPLOYMENT_NAME_MAP": "gpt-4o=prod-4o , other=x", + }) + p, err = ResolveSpecialProvider(spec, "gpt-4o", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://myres.openai.azure.com/openai/deployments/prod-4o"; got != want { + t.Fatalf("azure deployment base_url = %q, want %q", got, want) + } + + // --base-url flag wins over env origin. + p, err = ResolveSpecialProvider(spec, "gpt-4o", "https://flag.example.com", envFrom(map[string]string{"AZURE_OPENAI_API_KEY": "k"})) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://flag.example.com/openai/v1"; got != want { + t.Fatalf("azure flag base_url = %q, want %q", got, want) + } +} + +func TestParseDeploymentMap(t *testing.T) { + m := parseDeploymentMap(" a=1, b = 2 ,,bad,c=,=d ") + if m["a"] != "1" || m["b"] != "2" { + t.Fatalf("parseDeploymentMap = %v, want a=1 b=2", m) + } + if _, ok := m["bad"]; ok { + t.Fatalf("expected 'bad' (no '=') to be skipped: %v", m) + } + if _, ok := m["c"]; ok { + t.Fatalf("expected 'c=' (empty value) to be skipped: %v", m) + } +} + +func TestResolveSpecialProvider_Bedrock(t *testing.T) { + spec := specFor(t, "amazon-bedrock") + + // No credentials at all → names the bearer token. + if _, err := ResolveSpecialProvider(spec, "claude", "", envFrom(nil)); err == nil || + !strings.Contains(err.Error(), "AWS_BEARER_TOKEN_BEDROCK") { + t.Fatalf("expected AWS_BEARER_TOKEN_BEDROCK error, got %v", err) + } + + // Only profile present → SigV4-unsupported error, still names bearer token. + env := envFrom(map[string]string{"AWS_PROFILE": "default"}) + if _, err := ResolveSpecialProvider(spec, "claude", "", env); err == nil || + !strings.Contains(err.Error(), "SigV4") || + !strings.Contains(err.Error(), "AWS_BEARER_TOKEN_BEDROCK") { + t.Fatalf("expected SigV4-unsupported error naming bearer token, got %v", err) + } + + // Only static keys present → SigV4-unsupported error. + env = envFrom(map[string]string{"AWS_ACCESS_KEY_ID": "id", "AWS_SECRET_ACCESS_KEY": "secret"}) + if _, err := ResolveSpecialProvider(spec, "claude", "", env); err == nil || + !strings.Contains(err.Error(), "SigV4") { + t.Fatalf("expected SigV4-unsupported error for static keys, got %v", err) + } + + // Bearer token + default region. + env = envFrom(map[string]string{"AWS_BEARER_TOKEN_BEDROCK": "tok"}) + p, err := ResolveSpecialProvider(spec, "claude", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://bedrock-runtime.us-east-1.amazonaws.com"; got != want { + t.Fatalf("bedrock base_url = %q, want %q", got, want) + } + + // Bearer token + explicit region. + env = envFrom(map[string]string{"AWS_BEARER_TOKEN_BEDROCK": "tok", "AWS_REGION": "eu-west-1"}) + p, err = ResolveSpecialProvider(spec, "claude", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://bedrock-runtime.eu-west-1.amazonaws.com"; got != want { + t.Fatalf("bedrock base_url = %q, want %q", got, want) + } +} + +func TestResolveSpecialProvider_GoogleVertex(t *testing.T) { + spec := specFor(t, "google-vertex") + + if _, err := ResolveSpecialProvider(spec, "gemini", "", envFrom(nil)); err == nil || + !strings.Contains(err.Error(), "GOOGLE_CLOUD_PROJECT") { + t.Fatalf("expected GOOGLE_CLOUD_PROJECT error, got %v", err) + } + + env := envFrom(map[string]string{"GOOGLE_CLOUD_PROJECT": "proj"}) + if _, err := ResolveSpecialProvider(spec, "gemini", "", env); err == nil || + !strings.Contains(err.Error(), "GOOGLE_CLOUD_LOCATION") { + t.Fatalf("expected GOOGLE_CLOUD_LOCATION error, got %v", err) + } + + env = envFrom(map[string]string{"GOOGLE_CLOUD_PROJECT": "proj", "GOOGLE_CLOUD_LOCATION": "us-central1"}) + if _, err := ResolveSpecialProvider(spec, "gemini", "", env); err == nil || + !strings.Contains(err.Error(), "GOOGLE_CLOUD_API_KEY") || + !strings.Contains(err.Error(), "GOOGLE_APPLICATION_CREDENTIALS") { + t.Fatalf("expected credentials error naming both sources, got %v", err) + } + + // Fully configured with API key. + env = envFrom(map[string]string{ + "GOOGLE_CLOUD_PROJECT": "proj", + "GOOGLE_CLOUD_LOCATION": "us-central1", + "GOOGLE_CLOUD_API_KEY": "k", + }) + p, err := ResolveSpecialProvider(spec, "gemini", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://us-central1-aiplatform.googleapis.com"; got != want { + t.Fatalf("vertex base_url = %q, want %q", got, want) + } + + // ADC credential source also satisfies. + env = envFrom(map[string]string{ + "GOOGLE_CLOUD_PROJECT": "proj", + "GOOGLE_CLOUD_LOCATION": "europe-west4", + "GOOGLE_APPLICATION_CREDENTIALS": "/path/to/adc.json", + }) + p, err = ResolveSpecialProvider(spec, "gemini", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got, want := baseURLOf(t, p), "https://europe-west4-aiplatform.googleapis.com"; got != want { + t.Fatalf("vertex ADC base_url = %q, want %q", got, want) + } +} + +func TestResolveSpecialProvider_CloudflareWorkersAI(t *testing.T) { + spec := specFor(t, "cloudflare-workers-ai") + + if _, err := ResolveSpecialProvider(spec, "m", "", envFrom(nil)); err == nil || + !strings.Contains(err.Error(), "CLOUDFLARE_API_KEY") { + t.Fatalf("expected CLOUDFLARE_API_KEY error, got %v", err) + } + + env := envFrom(map[string]string{"CLOUDFLARE_API_KEY": "k"}) + if _, err := ResolveSpecialProvider(spec, "m", "", env); err == nil || + !strings.Contains(err.Error(), "CLOUDFLARE_ACCOUNT_ID") { + t.Fatalf("expected CLOUDFLARE_ACCOUNT_ID error, got %v", err) + } + + env = envFrom(map[string]string{"CLOUDFLARE_API_KEY": "k", "CLOUDFLARE_ACCOUNT_ID": "acct123"}) + p, err := ResolveSpecialProvider(spec, "m", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + want := "https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1" + if got := baseURLOf(t, p); got != want { + t.Fatalf("workers-ai base_url = %q, want %q", got, want) + } + if _, ok := p.(*openAICompatDriver); !ok { + t.Fatalf("workers-ai should speak OpenAI wire, got %T", p) + } +} + +func TestResolveSpecialProvider_CloudflareAIGateway(t *testing.T) { + spec := specFor(t, "cloudflare-ai-gateway") + + if _, err := ResolveSpecialProvider(spec, "m", "", envFrom(nil)); err == nil || + !strings.Contains(err.Error(), "CLOUDFLARE_API_KEY") { + t.Fatalf("expected CLOUDFLARE_API_KEY error, got %v", err) + } + + env := envFrom(map[string]string{"CLOUDFLARE_API_KEY": "k", "CLOUDFLARE_ACCOUNT_ID": "acct123"}) + if _, err := ResolveSpecialProvider(spec, "m", "", env); err == nil || + !strings.Contains(err.Error(), "CLOUDFLARE_GATEWAY_ID") { + t.Fatalf("expected CLOUDFLARE_GATEWAY_ID error, got %v", err) + } + + env = envFrom(map[string]string{ + "CLOUDFLARE_API_KEY": "k", + "CLOUDFLARE_ACCOUNT_ID": "acct123", + "CLOUDFLARE_GATEWAY_ID": "gw456", + }) + p, err := ResolveSpecialProvider(spec, "m", "", env) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + want := "https://gateway.ai.cloudflare.com/v1/acct123/gw456/anthropic" + if got := baseURLOf(t, p); got != want { + t.Fatalf("ai-gateway base_url = %q, want %q", got, want) + } + if _, ok := p.(*anthropicCompatDriver); !ok { + t.Fatalf("ai-gateway should speak Anthropic wire, got %T", p) + } +} + +func TestIsSpecialAuthProvider(t *testing.T) { + special := []string{ + "azure-openai-responses", "amazon-bedrock", "google-vertex", + "cloudflare-workers-ai", "cloudflare-ai-gateway", + } + for _, name := range special { + if !IsSpecialAuthProvider(specFor(t, name)) { + t.Errorf("%s should be a special-auth provider", name) + } + } + for _, name := range []string{"openai", "anthropic", "deepseek"} { + if IsSpecialAuthProvider(specFor(t, name)) { + t.Errorf("%s should NOT be a special-auth provider", name) + } + } +} diff --git a/pigo/internal/provider/thinking_test.go b/pigo/internal/provider/thinking_test.go new file mode 100644 index 0000000..df15520 --- /dev/null +++ b/pigo/internal/provider/thinking_test.go @@ -0,0 +1,264 @@ +// Tests for reasoning/thinking wiring across both wire protocols (issues +// #240-#244): request encoders forwarding ThinkingLevel, the Anthropic +// thinking-block echo on multi-turn tool-use messages, OpenAI reasoning_content +// stream decoding, the max_tokens fallback, and strict-gateway empty-content +// handling. +package provider + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// decodeBody unmarshals an encoded request body into a generic map for asserts. +func decodeBody(t *testing.T, b []byte) map[string]any { + t.Helper() + var m map[string]any + if err := json.Unmarshal(b, &m); err != nil { + t.Fatalf("unmarshal body: %v", err) + } + return m +} + +// TestOpenAIReasoningEffortEncoding verifies ThinkingLevel maps to +// reasoning_effort, and that off/unset omits the field entirely (#240). +func TestOpenAIReasoningEffortEncoding(t *testing.T) { + base := CompletionRequest{Model: "o3-mini", Context: LlmContext{}} + + // off / unset → no reasoning_effort. + for _, lvl := range []agentcore.ThinkingLevel{"", agentcore.ThinkingOff} { + base.Config.ThinkingLevel = lvl + b, err := encodeOpenAIRequest(base) + if err != nil { + t.Fatalf("encode: %v", err) + } + if _, ok := decodeBody(t, b)["reasoning_effort"]; ok { + t.Errorf("level %q: reasoning_effort should be omitted", lvl) + } + } + + cases := map[agentcore.ThinkingLevel]string{ + agentcore.ThinkingMinimal: "minimal", + agentcore.ThinkingLow: "low", + agentcore.ThinkingMedium: "medium", + agentcore.ThinkingHigh: "high", + agentcore.ThinkingXHigh: "high", + } + for lvl, want := range cases { + base.Config.ThinkingLevel = lvl + b, err := encodeOpenAIRequest(base) + if err != nil { + t.Fatalf("encode: %v", err) + } + if got := decodeBody(t, b)["reasoning_effort"]; got != want { + t.Errorf("level %q: reasoning_effort = %v, want %q", lvl, got, want) + } + } +} + +// TestAnthropicThinkingEncoding verifies ThinkingLevel enables the thinking +// block with a budget, and off/unset omits it (#240). +func TestAnthropicThinkingEncoding(t *testing.T) { + base := CompletionRequest{Model: "claude-x", Context: LlmContext{}} + + base.Config.ThinkingLevel = agentcore.ThinkingOff + b, _ := encodeAnthropicRequest(base, nil) + if _, ok := decodeBody(t, b)["thinking"]; ok { + t.Error("off: thinking block should be omitted") + } + + base.Config.ThinkingLevel = agentcore.ThinkingMedium + b, _ = encodeAnthropicRequest(base, nil) + th, ok := decodeBody(t, b)["thinking"].(map[string]any) + if !ok { + t.Fatal("medium: thinking block missing") + } + if th["type"] != "enabled" { + t.Errorf("thinking.type = %v, want enabled", th["type"]) + } + if bt, _ := th["budget_tokens"].(float64); bt <= 0 { + t.Errorf("thinking.budget_tokens = %v, want > 0", th["budget_tokens"]) + } +} + +// TestAnthropicMaxTokensFallback verifies the fallback prefers the model's +// MaxOutputTokens and otherwise uses the coding-friendly 8192 default (#243). +func TestAnthropicMaxTokensFallback(t *testing.T) { + req := CompletionRequest{Model: "claude-x", Context: LlmContext{}} + + // No model metadata → 8192 default. + b, _ := encodeAnthropicRequest(req, nil) + if got := decodeBody(t, b)["max_tokens"].(float64); got != 8192 { + t.Errorf("default max_tokens = %v, want 8192", got) + } + + // Model with a declared cap → that cap. + models := []Model{{ID: "claude-x", MaxOutputTokens: 12000}} + b, _ = encodeAnthropicRequest(req, models) + if got := decodeBody(t, b)["max_tokens"].(float64); got != 12000 { + t.Errorf("model-cap max_tokens = %v, want 12000", got) + } + + // Explicit Extra hint still wins. + req.Config.Extra = map[string]any{"max_tokens": 2000} + b, _ = encodeAnthropicRequest(req, models) + if got := decodeBody(t, b)["max_tokens"].(float64); got != 2000 { + t.Errorf("explicit max_tokens = %v, want 2000", got) + } +} + +// TestAnthropicThinkingBlockEcho verifies a multi-turn assistant message with a +// thinking block + tool call re-emits the thinking block (with signature) ahead +// of the tool_use block (#241). +func TestAnthropicThinkingBlockEcho(t *testing.T) { + think := agentcore.NewThinkingContent("let me reason") + think.ThinkingSignature = "sig-abc" + msg := agentcore.AssistantMessage{ + Content: agentcore.ContentList{ + think, + agentcore.NewToolCallContent("call_1", "read", json.RawMessage(`{"path":"x"}`)), + }, + } + entry := encodeAnthropicMessage(msg) + blocks, ok := entry["content"].([]map[string]any) + if !ok || len(blocks) != 2 { + t.Fatalf("want 2 content blocks, got %#v", entry["content"]) + } + if blocks[0]["type"] != "thinking" { + t.Errorf("first block type = %v, want thinking", blocks[0]["type"]) + } + if blocks[0]["signature"] != "sig-abc" { + t.Errorf("thinking signature = %v, want sig-abc", blocks[0]["signature"]) + } + if blocks[1]["type"] != "tool_use" { + t.Errorf("second block type = %v, want tool_use", blocks[1]["type"]) + } +} + +// TestAnthropicRedactedThinkingEcho verifies redacted thinking round-trips as a +// redacted_thinking block carrying the signature as data (#241). +func TestAnthropicRedactedThinkingEcho(t *testing.T) { + think := agentcore.ThinkingContent{Type: agentcore.ContentTypeThinking, Redacted: true, ThinkingSignature: "redacted-data"} + msg := agentcore.AssistantMessage{Content: agentcore.ContentList{think}} + entry := encodeAnthropicMessage(msg) + blocks := entry["content"].([]map[string]any) + if blocks[0]["type"] != "redacted_thinking" || blocks[0]["data"] != "redacted-data" { + t.Errorf("redacted block = %#v", blocks[0]) + } +} + +// TestOpenAIAssistantContentNullWithToolCalls verifies a tool-call-only +// assistant turn sends content:null (not ""), while a text-only turn keeps its +// text (#244). +func TestOpenAIAssistantContentNullWithToolCalls(t *testing.T) { + toolOnly := agentcore.AssistantMessage{ + Content: agentcore.ContentList{ + agentcore.NewToolCallContent("call_1", "read", json.RawMessage(`{}`)), + }, + } + entry := encodeOpenAIMessage(toolOnly)[0] + if entry["content"] != nil { + t.Errorf("tool-only content = %#v, want nil", entry["content"]) + } + if _, ok := entry["tool_calls"]; !ok { + t.Error("tool_calls missing") + } + + textOnly := agentcore.AssistantMessage{ + Content: agentcore.ContentList{agentcore.NewTextContent("hi")}, + } + if got := encodeOpenAIMessage(textOnly)[0]["content"]; got != "hi" { + t.Errorf("text content = %#v, want hi", got) + } +} + +// TestAnthropicEmptyAssistantNoEmptyText verifies an assistant message with no +// usable content does not emit an empty-string text block (#244). +func TestAnthropicEmptyAssistantNoEmptyText(t *testing.T) { + msg := agentcore.AssistantMessage{Content: agentcore.ContentList{agentcore.NewTextContent("")}} + entry := encodeAnthropicMessage(msg) + blocks := entry["content"].([]map[string]any) + if len(blocks) != 1 { + t.Fatalf("want 1 fallback block, got %d", len(blocks)) + } + if txt, _ := blocks[0]["text"].(string); strings.TrimSpace(txt) == "" && txt == "" { + t.Errorf("fallback text block is empty string, want non-empty placeholder") + } +} + +// TestAnthropicThinkingRaisesMaxTokens verifies max_tokens is lifted above the +// thinking budget: Anthropic requires budget_tokens < max_tokens, so a low cap +// must be raised to leave headroom for the visible reply (#243). +func TestAnthropicThinkingRaisesMaxTokens(t *testing.T) { + req := CompletionRequest{Model: "claude-x", Context: LlmContext{}} + req.Config.ThinkingLevel = agentcore.ThinkingXHigh // budget 32768 + + // Default cap (8192) is below the budget → must be raised above it. + b, _ := encodeAnthropicRequest(req, nil) + body := decodeBody(t, b) + budget := body["thinking"].(map[string]any)["budget_tokens"].(float64) + maxTok := body["max_tokens"].(float64) + if maxTok <= budget { + t.Errorf("max_tokens = %v, want > budget_tokens %v", maxTok, budget) + } + + // A caller cap already above budget+headroom is left untouched. + req.Config.Extra = map[string]any{"max_tokens": 100000} + b, _ = encodeAnthropicRequest(req, nil) + if got := decodeBody(t, b)["max_tokens"].(float64); got != 100000 { + t.Errorf("max_tokens = %v, want caller value 100000", got) + } +} + +// TestOpenAIReasoningContentDecoding verifies the decoder accumulates +// reasoning_content into a ThinkingContent block ahead of text (#242). +func TestOpenAIReasoningContentDecoding(t *testing.T) { + d := NewOpenAIDecoder() + chunks := []string{ + `{"id":"c1","choices":[{"delta":{"reasoning_content":"think "}}]}`, + `{"choices":[{"delta":{"reasoning_content":"harder"}}]}`, + `{"choices":[{"delta":{"content":"answer"}}]}`, + `{"choices":[{"finish_reason":"stop"}]}`, + } + var events []StreamEvent + for _, c := range chunks { + evs, err := d.Decode([]byte(c)) + if err != nil { + t.Fatalf("decode: %v", err) + } + events = append(events, evs...) + } + done, _ := d.Finish() + events = append(events, done...) + + var final agentcore.AssistantMessage + for _, e := range events { + if de, ok := e.(StreamDoneEvent); ok { + final = de.Message + } + } + if len(final.Content) != 2 { + t.Fatalf("want thinking+text, got %d blocks: %#v", len(final.Content), final.Content) + } + th, ok := final.Content[0].(agentcore.ThinkingContent) + if !ok || th.Thinking != "think harder" { + t.Errorf("block[0] = %#v, want thinking 'think harder'", final.Content[0]) + } + txt, ok := final.Content[1].(agentcore.TextContent) + if !ok || txt.Text != "answer" { + t.Errorf("block[1] = %#v, want text 'answer'", final.Content[1]) + } + + sawThinking := false + for _, e := range events { + if _, ok := e.(StreamThinkingEvent); ok { + sawThinking = true + } + } + if !sawThinking { + t.Error("expected at least one StreamThinkingEvent") + } +} diff --git a/pigo/internal/provider/transport.go b/pigo/internal/provider/transport.go new file mode 100644 index 0000000..8f3ea91 --- /dev/null +++ b/pigo/internal/provider/transport.go @@ -0,0 +1,393 @@ +// This file implements the shared transport driver (US-007 support): a +// provider-agnostic layer that turns an HTTP request into a stream of +// StreamEvents. Each provider degenerates to a stateful Decoder; the transport +// owns HTTP + SSE line parsing + retry + dual watchdogs + the dual failure +// model. +// +// The design mirrors pi's providerio: the transport never returns a runtime +// failure as a Go error once streaming has begun — it rides the stream as a +// terminal StreamErrorEvent. Only the earliest "cannot build the stream" case +// (bad request construction) is a returned error. +package provider + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "strconv" + "strings" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// StreamEvent is the transport-level alias for AssistantMessageEvent. Decoders +// produce these; the transport forwards them onto the stream. (Decision #25: +// reuse AssistantMessageEvent rather than a parallel event type.) +type StreamEvent = AssistantMessageEvent + +// Decoder is the per-provider stateful SSE payload decoder. The transport calls +// Decode for every complete SSE data payload (one event's worth of bytes) and +// Finish once the stream ends so the decoder can flush any buffered terminal +// event. +type Decoder interface { + // Decode turns one SSE data payload into zero or more StreamEvents. A + // returned error is treated as a runtime stream failure (terminal error + // event), never a panic. + Decode(payload []byte) ([]StreamEvent, error) + // Finish flushes any trailing state, returning a final batch of events. + Finish() ([]StreamEvent, error) +} + +// defaultIdleTimeout is the watchdog idle window; PIGO_STREAM_IDLE_TIMEOUT +// (a Go duration string, e.g. "3m") overrides it. +const defaultIdleTimeout = 5 * time.Minute + +const ( + // defaultMaxConnectRetries bounds connect-only retries when TransportConfig + // leaves MaxConnectRetries at zero. + defaultMaxConnectRetries = 2 + // statusTooManyRequestsCF (529) is Cloudflare's "site overloaded" status, + // which some upstreams also emit; treated as retryable alongside 429/503. + statusTooManyRequestsCF = 529 + // stallFactor slackens the content-stall watchdog relative to the idle + // window (stall = idle × stallFactor) so a slow-but-progressing stream is not + // killed by the stall guard. + stallFactor = 1.2 + // errorBodyLimit bounds how many bytes of an upstream error body are read + // into the returned error message. + errorBodyLimit = 4096 +) + +// idleTimeout resolves the configured idle watchdog window. +func idleTimeout() time.Duration { + if v := os.Getenv("PIGO_STREAM_IDLE_TIMEOUT"); v != "" { + if d, err := time.ParseDuration(v); err == nil && d > 0 { + return d + } + } + return defaultIdleTimeout +} + +// TransportConfig configures a single StreamRequest run. +type TransportConfig struct { + // Client is the HTTP client; defaults to http.DefaultClient when nil. + Client *http.Client + // NewRequest builds a fresh *http.Request for each connection attempt. It is + // called once per connect (initial + reconnects) so retries never replay a + // consumed body — the caller owns idempotent request construction. + NewRequest func(ctx context.Context) (*http.Request, error) + // Decoder converts SSE payloads to StreamEvents (required). + Decoder Decoder + // MaxConnectRetries bounds connect-only retries (default 2). + MaxConnectRetries int +} + +// StreamRequest runs cfg as a transport stream. Per the dual failure model it +// returns an error only when the very first request cannot be built or the +// initial connection can never be established; every runtime failure once +// streaming begins rides the returned stream as a terminal StreamErrorEvent. +func StreamRequest(ctx context.Context, cfg TransportConfig) (*AssistantMessageEventStream, error) { + if cfg.NewRequest == nil { + return nil, errors.New("transport: NewRequest is required") + } + if cfg.Decoder == nil { + return nil, errors.New("transport: Decoder is required") + } + client := cfg.Client + if client == nil { + client = http.DefaultClient + } + maxRetries := cfg.MaxConnectRetries + if maxRetries == 0 { + maxRetries = defaultMaxConnectRetries + } + + // Connect once up front so a "cannot even build the stream" failure surfaces + // as a returned error (the only early-error case per FR-13). + resp, err := connect(ctx, client, cfg.NewRequest, maxRetries) + if err != nil { + return nil, err + } + + stream := NewAssistantMessageEventStream(0) + go pump(ctx, stream, resp, cfg.Decoder) + return stream, nil +} + +// connect performs the initial connection with retry. It only retries when the +// server explicitly signals a retryable condition (429/503/529); it never +// replays a consumed stream, so retrying at connect time is always safe. +func connect(ctx context.Context, client *http.Client, newReq func(context.Context) (*http.Request, error), maxRetries int) (*http.Response, error) { + var lastErr error + for attempt := 0; attempt <= maxRetries; attempt++ { + req, err := newReq(ctx) + if err != nil { + return nil, fmt.Errorf("transport: build request: %w", err) + } + resp, err := client.Do(req) + if err != nil { + lastErr = classifyTransportError(err) + if !isRetryableNetErr(err) || attempt == maxRetries { + return nil, lastErr + } + if !sleepBackoff(ctx, attempt, 0) { + return nil, ctx.Err() + } + continue + } + if resp.StatusCode == http.StatusTooManyRequests || + resp.StatusCode == http.StatusServiceUnavailable || + resp.StatusCode == statusTooManyRequestsCF { + wait := retryAfter(resp.Header) + resp.Body.Close() + lastErr = fmt.Errorf("transport: upstream %d", resp.StatusCode) + if attempt == maxRetries { + return nil, lastErr + } + if !sleepBackoff(ctx, attempt, wait) { + return nil, ctx.Err() + } + continue + } + if resp.StatusCode >= 400 { + body, _ := io.ReadAll(io.LimitReader(resp.Body, errorBodyLimit)) + resp.Body.Close() + return nil, fmt.Errorf("transport: upstream %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + return resp, nil + } + return nil, lastErr +} + +// pump drives the SSE read loop with dual watchdogs, decoding payloads and +// forwarding events onto the stream. All runtime failures become a terminal +// error event; pump always closes the stream. +func pump(ctx context.Context, stream *AssistantMessageEventStream, resp *http.Response, dec Decoder) { + defer stream.Close() + defer resp.Body.Close() + + idle := idleTimeout() + // content-stall watchdog is slightly slacker than idle (idle × stallFactor) + // so a slow but progressing stream is not killed by the stall guard. + stall := time.Duration(float64(idle) * stallFactor) + + // The watchdog fires by cancelling a derived context; reads race against it. + watchCtx, cancel := context.WithCancel(ctx) + defer cancel() + + // done stops the reader goroutine so it never blocks on a send after pump + // returns (watchdog / abort paths), avoiding a goroutine leak. + done := make(chan struct{}) + defer close(done) + lines := make(chan string) + readErr := make(chan error, 1) + go readLines(resp.Body, lines, readErr, done) + + var dataBuf strings.Builder + idleTimer := time.NewTimer(idle) + stallTimer := time.NewTimer(stall) + defer idleTimer.Stop() + defer stallTimer.Stop() + + emit := func(events []StreamEvent) bool { + for _, ev := range events { + if err := stream.Emit(watchCtx, ev); err != nil { + return false + } + } + return true + } + + fail := func(msg string, err error) { + stream.Emit(context.Background(), StreamErrorEvent{ + Message: agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + StopReason: agentcore.StopReasonError, + ErrorMessage: msg, + }, + Err: err, + }) + } + + flush := func() bool { + if dataBuf.Len() == 0 { + return true + } + payload := dataBuf.String() + dataBuf.Reset() + if payload == "[DONE]" { + return true + } + events, err := dec.Decode([]byte(payload)) + if err != nil { + fail("decode error: "+err.Error(), err) + return false + } + return emit(events) + } + + for { + select { + case <-ctx.Done(): + fail("stream aborted", ctx.Err()) + return + case <-idleTimer.C: + fail("idle timeout: no data received", errStreamIdle) + return + case <-stallTimer.C: + fail("content stall timeout", errStreamStall) + return + case err := <-readErr: + if err != nil && !errors.Is(err, io.EOF) { + fail("read error: "+classifyTransportError(err).Error(), err) + return + } + // Clean EOF: flush any buffered payload, then finish the decoder. + if !flush() { + return + } + finalEvents, ferr := dec.Finish() + if ferr != nil { + fail("finish error: "+ferr.Error(), ferr) + return + } + emit(finalEvents) + return + case line, ok := <-lines: + if !ok { + continue + } + // Any byte resets the idle watchdog; a flushed event resets stall. + resetTimer(idleTimer, idle) + line = strings.TrimRight(line, "\r\n") + switch { + case line == "": + // Blank line = event boundary: flush accumulated data. + if !flush() { + return + } + resetTimer(stallTimer, stall) + case strings.HasPrefix(line, ":"): + // Comment / keep-alive: ignore payload, watchdog already reset. + case strings.HasPrefix(line, "data:"): + dataBuf.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + default: + // Non-data field (event:, id:, etc.) — ignored for our decoders. + } + } + } +} + +// readLines reads the body line by line, pushing each onto lines and the final +// error (io.EOF on clean close) onto readErr. It stops promptly when done is +// closed so pump can return on a watchdog/abort without leaking this goroutine. +func readLines(r io.Reader, lines chan<- string, readErr chan<- error, done <-chan struct{}) { + br := bufio.NewReader(r) + for { + line, err := br.ReadString('\n') + if line != "" { + select { + case lines <- line: + case <-done: + return + } + } + if err != nil { + select { + case readErr <- err: + case <-done: + } + return + } + } +} + +// resetTimer stops and re-arms t to fire after d. +func resetTimer(t *time.Timer, d time.Duration) { + if !t.Stop() { + select { + case <-t.C: + default: + } + } + t.Reset(d) +} + +// Sentinel errors for watchdog classification. +var ( + errStreamIdle = errors.New("stream idle timeout") + errStreamStall = errors.New("stream content stall") +) + +// retryAfter parses a Retry-After header (seconds or HTTP-date), returning 0 +// when absent/unparseable. +func retryAfter(h http.Header) time.Duration { + v := h.Get("Retry-After") + if v == "" { + return 0 + } + if secs, err := strconv.Atoi(v); err == nil && secs >= 0 { + return time.Duration(secs) * time.Second + } + if t, err := http.ParseTime(v); err == nil { + if d := time.Until(t); d > 0 { + return d + } + } + return 0 +} + +// sleepBackoff waits for the retry delay (Retry-After if given, else +// exponential), honoring ctx cancellation. Returns false if ctx was cancelled. +func sleepBackoff(ctx context.Context, attempt int, retryAfter time.Duration) bool { + d := retryAfter + if d == 0 { + d = time.Duration(1< 0 { + i-- + buf[i] = byte('0' + n%10) + n /= 10 + } + return string(buf[i:]) +} diff --git a/pigo/internal/remotecontrol/bridge_test.go b/pigo/internal/remotecontrol/bridge_test.go new file mode 100644 index 0000000..79234ec --- /dev/null +++ b/pigo/internal/remotecontrol/bridge_test.go @@ -0,0 +1,192 @@ +package remotecontrol + +import ( + "context" + "sync" + "testing" + "time" +) + +type fakeSink struct { + mu sync.Mutex + outputs []string + confirms []confirmReq + connected bool +} + +type confirmReq struct { + id, tool, summary string +} + +func (f *fakeSink) SendOutput(text string) { + f.mu.Lock() + defer f.mu.Unlock() + f.outputs = append(f.outputs, text) +} + +func (f *fakeSink) SendConfirm(id, tool, summary string) { + f.mu.Lock() + defer f.mu.Unlock() + f.confirms = append(f.confirms, confirmReq{id, tool, summary}) +} + +func (f *fakeSink) HasClient() bool { + f.mu.Lock() + defer f.mu.Unlock() + return f.connected +} + +func (f *fakeSink) lastConfirmID() string { + f.mu.Lock() + defer f.mu.Unlock() + if len(f.confirms) == 0 { + return "" + } + return f.confirms[len(f.confirms)-1].id +} + +func TestBridgeOutputWriterTees(t *testing.T) { + sink := &fakeSink{} + b := NewBridge(sink) + w := b.OutputWriter() + n, err := w.Write([]byte("hello world")) + if err != nil || n != len("hello world") { + t.Fatalf("Write = (%d,%v), want (%d,nil)", n, err, len("hello world")) + } + sink.mu.Lock() + defer sink.mu.Unlock() + if len(sink.outputs) != 1 || sink.outputs[0] != "hello world" { + t.Fatalf("outputs = %v, want [hello world]", sink.outputs) + } +} + +func TestBridgeRemoteInput(t *testing.T) { + b := NewBridge(&fakeSink{}) + b.OnInput("do a thing") + select { + case got := <-b.RemoteInput(): + if got != "do a thing" { + t.Fatalf("input = %q, want 'do a thing'", got) + } + case <-time.After(time.Second): + t.Fatal("no input delivered") + } +} + +func TestBridgeOnInputDropsOldestWhenFull(t *testing.T) { + b := NewBridge(&fakeSink{}) + // Fill beyond capacity; must not block and must retain the newest items. + total := remoteInputBuffer + 10 + for i := range total { + b.OnInput(itoa(uint64(i))) + } + // Drain and ensure we get exactly buffer-size items, ending at the newest. + var got []string + for { + select { + case v := <-b.RemoteInput(): + got = append(got, v) + continue + default: + } + break + } + if len(got) != remoteInputBuffer { + t.Fatalf("drained %d items, want %d", len(got), remoteInputBuffer) + } + if last := got[len(got)-1]; last != itoa(uint64(total-1)) { + t.Fatalf("newest = %q, want %q", last, itoa(uint64(total-1))) + } +} + +func TestBridgeConfirmResolvedRemotely(t *testing.T) { + sink := &fakeSink{connected: true} + b := NewBridge(sink) + + done := make(chan struct { + d Decision + remote bool + }, 1) + go func() { + d, remote := b.Confirm(context.Background(), "shell", "rm -rf /tmp/x") + done <- struct { + d Decision + remote bool + }{d, remote} + }() + + // Wait for the confirm to be sent, then answer it via OnDecide. + deadline := time.Now().Add(time.Second) + for sink.lastConfirmID() == "" { + if time.Now().After(deadline) { + t.Fatal("confirm never sent") + } + time.Sleep(2 * time.Millisecond) + } + b.OnDecide(sink.lastConfirmID(), true, true) + + select { + case r := <-done: + if !r.remote { + t.Fatal("remote flag = false, want true") + } + if !r.d.Approve || !r.d.Always { + t.Fatalf("decision = %+v, want approve+always", r.d) + } + case <-time.After(time.Second): + t.Fatal("Confirm did not return") + } +} + +func TestBridgeConfirmCancelledByContext(t *testing.T) { + sink := &fakeSink{connected: true} + b := NewBridge(sink) + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan bool, 1) + go func() { + _, remote := b.Confirm(ctx, "shell", "ls") + done <- remote + }() + // Let the confirm register, then cancel (simulating a local answer). + deadline := time.Now().Add(time.Second) + for sink.lastConfirmID() == "" { + if time.Now().After(deadline) { + t.Fatal("confirm never sent") + } + time.Sleep(2 * time.Millisecond) + } + cancel() + + select { + case remote := <-done: + if remote { + t.Fatal("remote flag = true, want false on ctx cancel") + } + case <-time.After(time.Second): + t.Fatal("Confirm did not return on cancel") + } +} + +func TestBridgeResolveConfirmUnknown(t *testing.T) { + b := NewBridge(&fakeSink{}) + if b.ResolveConfirm("nope", true, false) { + t.Fatal("ResolveConfirm on unknown id = true, want false") + } +} + +func TestBridgeEnabled(t *testing.T) { + var nilBridge *Bridge + if nilBridge.Enabled() { + t.Fatal("nil bridge Enabled = true, want false") + } + sink := &fakeSink{connected: false} + b := NewBridge(sink) + if b.Enabled() { + t.Fatal("Enabled = true with no client, want false") + } + sink.connected = true + if !b.Enabled() { + t.Fatal("Enabled = false with client, want true") + } +} diff --git a/pigo/internal/remotecontrol/lanaddr.go b/pigo/internal/remotecontrol/lanaddr.go new file mode 100644 index 0000000..6a70c76 --- /dev/null +++ b/pigo/internal/remotecontrol/lanaddr.go @@ -0,0 +1,116 @@ +package remotecontrol + +import ( + "errors" + "fmt" + "net" +) + +// ErrNoLAN is returned by DetectRoutableIP when no routable, non-loopback IPv4 +// address can be found. Binding loopback would produce a URL the phone cannot +// reach, so the caller must surface this rather than fall back silently. +var ErrNoLAN = errors.New("remotecontrol: no routable LAN address found; are you connected to Wi-Fi?") + +// ifaceInfo pairs a network interface's flags with its addresses. DetectRoutableIP +// consumes these so it can skip down / loopback / point-to-point (VPN tunnel) +// interfaces before considering their addresses. +type ifaceInfo struct { + flags net.Flags + addrs []net.Addr +} + +// ifaceLister returns the host's interfaces with their addresses. It is a package +// variable so tests can inject a fake set without real hardware. +var ifaceLister = defaultIfaceLister + +func defaultIfaceLister() ([]ifaceInfo, error) { + ifaces, err := net.Interfaces() + if err != nil { + return nil, err + } + out := make([]ifaceInfo, 0, len(ifaces)) + for _, ifi := range ifaces { + addrs, err := ifi.Addrs() + if err != nil { + continue // an interface whose addresses can't be read is unusable + } + out = append(out, ifaceInfo{flags: ifi.Flags, addrs: addrs}) + } + return out, nil +} + +// DetectRoutableIP returns a routable IPv4 address suitable for embedding in the +// printed pairing URL — one a phone on the same Wi-Fi can actually reach. +// +// It iterates interfaces (not bare addresses) so it can skip the ones that would +// yield an unreachable URL: interfaces that are down, loopback, or point-to-point +// (VPN utun / tunnel links, whose address the phone cannot route to). Among the +// rest it prefers a private LAN address (RFC1918: 10/8, 172.16/12, 192.168/16), +// which is what home/office Wi-Fi hands out; a non-private routable address is +// used only as a fallback when no private one exists. Callers that need a specific +// interface should override via Config.Host. It returns ErrNoLAN when nothing +// routable exists. +func DetectRoutableIP() (string, error) { + ifaces, err := ifaceLister() + if err != nil { + return "", fmt.Errorf("remotecontrol: list interfaces: %w", err) + } + var fallback string + for _, ifi := range ifaces { + if ifi.flags&net.FlagUp == 0 { + continue // interface is down + } + if ifi.flags&net.FlagLoopback != 0 { + continue + } + if ifi.flags&net.FlagPointToPoint != 0 { + continue // VPN / tunnel link — its address isn't reachable from the LAN + } + for _, a := range ifi.addrs { + var ip net.IP + switch v := a.(type) { + case *net.IPNet: + ip = v.IP + case *net.IPAddr: + ip = v.IP + default: + continue + } + v4 := ip.To4() + if v4 == nil { + continue // skip IPv6 for the LAN URL + } + if v4.IsLoopback() || v4.IsLinkLocalUnicast() || v4.IsLinkLocalMulticast() || v4.IsUnspecified() { + continue + } + if v4.IsPrivate() { + return v4.String(), nil + } + if fallback == "" { + fallback = v4.String() + } + } + } + if fallback != "" { + return fallback, nil + } + return "", ErrNoLAN +} + +// ListenFreePort binds a TCP listener on host. It first tries the requested +// port; if port is 0 or already in use it falls back to a kernel-assigned free +// port (:0). It returns the listener and the actual port bound. +func ListenFreePort(host string, port int) (net.Listener, int, error) { + if port != 0 { + ln, err := net.Listen("tcp", net.JoinHostPort(host, fmt.Sprint(port))) + if err == nil { + return ln, ln.Addr().(*net.TCPAddr).Port, nil + } + // Requested port unavailable — fall through to an auto-assigned one. + } + ln, err := net.Listen("tcp", net.JoinHostPort(host, "0")) + if err != nil { + return nil, 0, fmt.Errorf("remotecontrol: bind %s: %w", host, err) + } + return ln, ln.Addr().(*net.TCPAddr).Port, nil +} diff --git a/pigo/internal/remotecontrol/lanaddr_test.go b/pigo/internal/remotecontrol/lanaddr_test.go new file mode 100644 index 0000000..ef228ec --- /dev/null +++ b/pigo/internal/remotecontrol/lanaddr_test.go @@ -0,0 +1,182 @@ +package remotecontrol + +import ( + "net" + "testing" +) + +// withIfaceLister swaps the package ifaceLister for the duration of a test. +func withIfaceLister(t *testing.T, fn func() ([]ifaceInfo, error)) { + t.Helper() + orig := ifaceLister + ifaceLister = fn + t.Cleanup(func() { ifaceLister = orig }) +} + +// ipnets wraps IPs as *net.IPNet addresses for an ifaceInfo. +func ipnets(ips ...string) []net.Addr { + out := make([]net.Addr, 0, len(ips)) + for _, s := range ips { + out = append(out, &net.IPNet{IP: net.ParseIP(s)}) + } + return out +} + +func TestDetectRoutableIPPicksRoutableV4(t *testing.T) { + withIfaceLister(t, func() ([]ifaceInfo, error) { + return []ifaceInfo{{ + flags: net.FlagUp | net.FlagBroadcast | net.FlagMulticast, + addrs: ipnets("fe80::1", "127.0.0.1", "169.254.1.5", "192.168.1.42", "10.0.0.9"), + }}, nil + }) + ip, err := DetectRoutableIP() + if err != nil { + t.Fatalf("DetectRoutableIP: %v", err) + } + if ip != "192.168.1.42" { + t.Fatalf("ip = %q, want 192.168.1.42 (first private v4)", ip) + } +} + +// A VPN utun tunnel (point-to-point) whose address sorts before Wi-Fi must be +// skipped so the QR URL points at the LAN address the phone can reach — the +// white-screen bug this replaces. +func TestDetectRoutableIPSkipsVPNPointToPoint(t *testing.T) { + withIfaceLister(t, func() ([]ifaceInfo, error) { + return []ifaceInfo{ + { // loopback + flags: net.FlagUp | net.FlagLoopback, + addrs: ipnets("127.0.0.1"), + }, + { // VPN tunnel — routable but unreachable from the LAN + flags: net.FlagUp | net.FlagPointToPoint | net.FlagMulticast, + addrs: ipnets("172.31.201.147"), + }, + { // Wi-Fi + flags: net.FlagUp | net.FlagBroadcast | net.FlagMulticast, + addrs: ipnets("192.168.3.54"), + }, + }, nil + }) + ip, err := DetectRoutableIP() + if err != nil { + t.Fatalf("DetectRoutableIP: %v", err) + } + if ip != "192.168.3.54" { + t.Fatalf("ip = %q, want 192.168.3.54 (Wi-Fi, VPN skipped)", ip) + } +} + +// A down interface must not be chosen even if it carries a private address. +func TestDetectRoutableIPSkipsDownInterface(t *testing.T) { + withIfaceLister(t, func() ([]ifaceInfo, error) { + return []ifaceInfo{ + { // down: skipped despite the private address + flags: net.FlagBroadcast | net.FlagMulticast, + addrs: ipnets("192.168.9.9"), + }, + { // up + flags: net.FlagUp | net.FlagBroadcast | net.FlagMulticast, + addrs: ipnets("10.1.2.3"), + }, + }, nil + }) + ip, err := DetectRoutableIP() + if err != nil { + t.Fatalf("DetectRoutableIP: %v", err) + } + if ip != "10.1.2.3" { + t.Fatalf("ip = %q, want 10.1.2.3 (down iface skipped)", ip) + } +} + +// With no private address, a routable public address is used as a fallback. +func TestDetectRoutableIPFallsBackToPublic(t *testing.T) { + withIfaceLister(t, func() ([]ifaceInfo, error) { + return []ifaceInfo{{ + flags: net.FlagUp | net.FlagBroadcast | net.FlagMulticast, + addrs: ipnets("203.0.113.7"), + }}, nil + }) + ip, err := DetectRoutableIP() + if err != nil { + t.Fatalf("DetectRoutableIP: %v", err) + } + if ip != "203.0.113.7" { + t.Fatalf("ip = %q, want 203.0.113.7 (public fallback)", ip) + } +} + +func TestDetectRoutableIPSkipsLoopbackAndLinkLocal(t *testing.T) { + withIfaceLister(t, func() ([]ifaceInfo, error) { + return []ifaceInfo{{ + flags: net.FlagUp | net.FlagMulticast, + addrs: ipnets("127.0.0.1", "169.254.1.5"), + }}, nil + }) + if _, err := DetectRoutableIP(); err != ErrNoLAN { + t.Fatalf("err = %v, want ErrNoLAN", err) + } +} + +func TestDetectRoutableIPNoInterfaces(t *testing.T) { + withIfaceLister(t, func() ([]ifaceInfo, error) { + return nil, nil + }) + if _, err := DetectRoutableIP(); err != ErrNoLAN { + t.Fatalf("err = %v, want ErrNoLAN", err) + } +} + +func TestListenFreePortAutoAssign(t *testing.T) { + ln, port, err := ListenFreePort("127.0.0.1", 0) + if err != nil { + t.Fatalf("ListenFreePort: %v", err) + } + defer ln.Close() + if port <= 0 { + t.Fatalf("port = %d, want > 0", port) + } + if got := ln.Addr().(*net.TCPAddr).Port; got != port { + t.Fatalf("listener port %d != returned %d", got, port) + } +} + +func TestListenFreePortFallsBackWhenOccupied(t *testing.T) { + // Occupy a port, then ask ListenFreePort for that same port; it must fall + // back to a different, free one instead of failing. + occupied, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("occupy: %v", err) + } + defer occupied.Close() + busyPort := occupied.Addr().(*net.TCPAddr).Port + + ln, port, err := ListenFreePort("127.0.0.1", busyPort) + if err != nil { + t.Fatalf("ListenFreePort: %v", err) + } + defer ln.Close() + if port == busyPort { + t.Fatalf("port = %d, expected fallback away from occupied %d", port, busyPort) + } +} + +func TestListenFreePortUsesRequestedWhenFree(t *testing.T) { + // Find a free port, release it, then request it explicitly. + probe, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("probe: %v", err) + } + want := probe.Addr().(*net.TCPAddr).Port + probe.Close() + + ln, port, err := ListenFreePort("127.0.0.1", want) + if err != nil { + t.Fatalf("ListenFreePort: %v", err) + } + defer ln.Close() + if port != want { + t.Fatalf("port = %d, want requested %d", port, want) + } +} diff --git a/pigo/internal/remotecontrol/protocol.go b/pigo/internal/remotecontrol/protocol.go new file mode 100644 index 0000000..5e42423 --- /dev/null +++ b/pigo/internal/remotecontrol/protocol.go @@ -0,0 +1,47 @@ +package remotecontrol + +// FrameType enumerates the WebSocket message kinds exchanged between the server +// and the paired browser. It is the wire contract shared with the embedded SPA. +type FrameType string + +const ( + // FrameOutput streams session text from server to client. + FrameOutput FrameType = "output" + // FrameInput carries a prompt line submitted by the client to the server. + FrameInput FrameType = "input" + // FrameConfirm asks the client to approve/reject a risky tool call. + FrameConfirm FrameType = "confirm" + // FrameDecide carries the client's approval decision back to the server. + FrameDecide FrameType = "decide" + // FrameStatus reports session lifecycle changes to the client. + FrameStatus FrameType = "status" +) + +// Status values carried in Frame.State for FrameStatus frames. +const ( + StatusConnected = "connected" + StatusEnded = "ended" + StatusDisconnected = "disconnected" +) + +// Frame is a single WebSocket message. Fields are populated according to Type; +// unused fields are omitted from the JSON encoding. +type Frame struct { + Type FrameType `json:"type"` + + // Output + Text string `json:"text,omitempty"` + + // Confirm + ConfirmID string `json:"confirmId,omitempty"` + Tool string `json:"tool,omitempty"` + Summary string `json:"summary,omitempty"` + + // Decide + Approve bool `json:"approve,omitempty"` + Always bool `json:"always,omitempty"` + + // Status + State string `json:"state,omitempty"` // connected | ended | disconnected + Reason string `json:"reason,omitempty"` // human-readable detail for State +} diff --git a/pigo/internal/remotecontrol/qr.go b/pigo/internal/remotecontrol/qr.go new file mode 100644 index 0000000..9b9e4fd --- /dev/null +++ b/pigo/internal/remotecontrol/qr.go @@ -0,0 +1,51 @@ +package remotecontrol + +import ( + "strings" + + qrcode "github.com/skip2/go-qrcode" +) + +// Render encodes url as a QR code drawn with Unicode half-block characters, +// suitable for scanning off a terminal by a phone camera. Two matrix rows are +// packed into each text line (▀ ▄ █ and space), halving the printed height. +// +// The rendering assumes a dark-background terminal: light QR modules are drawn +// as bright block glyphs and dark modules as the terminal background. go-qrcode +// includes the mandatory quiet-zone border in its bitmap. +// +// On any encoding error Render returns ("", err); callers should degrade +// gracefully by printing the URL alone (the QR is a convenience, not required). +func Render(url string) (string, error) { + q, err := qrcode.New(url, qrcode.Medium) + if err != nil { + return "", err + } + bm := q.Bitmap() // true = dark module; includes quiet-zone border. + + var b strings.Builder + for y := 0; y < len(bm); y += 2 { + row := bm[y] + for x := 0; x < len(row); x++ { + // A "light" pixel is drawn as a block; a "dark" pixel is left as + // background. Rows beyond the matrix are treated as light (quiet). + topLight := !bm[y][x] + botLight := true + if y+1 < len(bm) { + botLight = !bm[y+1][x] + } + switch { + case topLight && botLight: + b.WriteRune('█') + case topLight && !botLight: + b.WriteRune('▀') + case !topLight && botLight: + b.WriteRune('▄') + default: + b.WriteByte(' ') + } + } + b.WriteByte('\n') + } + return b.String(), nil +} diff --git a/pigo/internal/remotecontrol/qr_test.go b/pigo/internal/remotecontrol/qr_test.go new file mode 100644 index 0000000..543366c --- /dev/null +++ b/pigo/internal/remotecontrol/qr_test.go @@ -0,0 +1,60 @@ +package remotecontrol + +import ( + "strings" + "testing" +) + +func TestRenderProducesBlockOutput(t *testing.T) { + out, err := Render("http://192.168.1.42:8080/pair?t=deadbeef") + if err != nil { + t.Fatalf("Render: %v", err) + } + if out == "" { + t.Fatal("Render returned empty output") + } + // Output must consist only of the half-block glyphs, spaces, and newlines. + for _, r := range out { + switch r { + case '█', '▀', '▄', ' ', '\n': + default: + t.Fatalf("unexpected rune %q in QR output", r) + } + } + // Must contain at least one dark-bearing glyph (not all blanks). + if !strings.ContainsAny(out, "▀▄ ") { + t.Fatal("QR output has no dark modules") + } +} + +func TestRenderIsDeterministic(t *testing.T) { + const url = "http://10.0.0.9:5000/pair?t=abc123" + a, err := Render(url) + if err != nil { + t.Fatalf("Render: %v", err) + } + b, err := Render(url) + if err != nil { + t.Fatalf("Render: %v", err) + } + if a != b { + t.Fatal("Render is not deterministic for the same URL") + } +} + +func TestRenderSquareRows(t *testing.T) { + out, err := Render("http://127.0.0.1:1/pair?t=x") + if err != nil { + t.Fatalf("Render: %v", err) + } + lines := strings.Split(strings.TrimRight(out, "\n"), "\n") + if len(lines) == 0 { + t.Fatal("no lines") + } + width := len([]rune(lines[0])) + for i, ln := range lines { + if got := len([]rune(ln)); got != width { + t.Fatalf("line %d width = %d, want uniform %d", i, got, width) + } + } +} diff --git a/pigo/internal/remotecontrol/server.go b/pigo/internal/remotecontrol/server.go new file mode 100644 index 0000000..40e598e --- /dev/null +++ b/pigo/internal/remotecontrol/server.go @@ -0,0 +1,602 @@ +package remotecontrol + +import ( + "context" + "embed" + "errors" + "fmt" + "io/fs" + "net" + "net/http" + "sync" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" +) + +// spaFiles holds the embedded browser SPA. The placeholder shipped here is +// fleshed out by the web-SPA node; the server only needs a valid FS to embed. +// +//go:embed web +var spaFiles embed.FS + +// Default configuration values (see tasks/spec-remote-control.md §3.2). +const ( + defaultPairTTL = 10 * time.Minute + cookieName = "pigo_rc" + maxFrameBytes = 64 * 1024 + + // outputFlushInterval is the coalescing tick: adjacent output writes buffered + // within one interval are sent as a single frame, so a fast stream produces a + // few batched frames rather than one per write, while still hitting the PRD's + // sub-100ms latency target (§8.2). + outputFlushInterval = 16 * time.Millisecond + // outputQueueMax bounds the pending buffer. A producer that would push it past + // this blocks until the pump drains, applying backpressure — the terminal + // mirror must not drop bytes (§5.4), so we never discard, only slow down. + outputQueueMax = 256 * 1024 + // replayRingBytes bounds the rolling output history kept for late-join / + // reconnect replay so a freshly connected browser sees recent context (§8.2). + replayRingBytes = 256 * 1024 +) + +// Config controls a remote-control server instance. +type Config struct { + PairTTL time.Duration // pairing-token lifetime; 0 → defaultPairTTL + Host string // LAN IP for the printed URL; "" → auto-detect + Port int // 0 → auto-pick, fall back on conflict + ConfirmTimeout time.Duration // 0 → wait forever for a remote decision + + // OnClientConnect, if set, is invoked (on the WebSocket goroutine) when a + // browser pairs and connects, with the client's remote address. The REPL uses + // it to print a one-line terminal notice so the operator notices remote access + // (§7.3). It must not block for long. + OnClientConnect func(remoteAddr string) + // OnClientDisconnect, if set, is invoked when the controlling client's + // WebSocket closes. It must not block for long. + OnClientDisconnect func() +} + +// Handler receives frames the browser sends. The REPL bridge implements it; +// tests supply a fake. Callbacks run on the WebSocket read goroutine and must +// not block for long. +type Handler interface { + // OnInput is called when the client submits a prompt line. + OnInput(text string) + // OnDecide is called when the client answers a confirmation request. + OnDecide(confirmID string, approve, always bool) +} + +// serverState tracks the lifecycle for gating requests. +type serverState int + +const ( + stateIdle serverState = iota + stateListening + stateEnded +) + +// Server is an in-process HTTP + WebSocket server that mirrors the CLI session +// to a single paired browser on the LAN. +type Server struct { + cfg Config + tokens *TokenStore + handler Handler + spa fs.FS + + mu sync.Mutex + state serverState + ln net.Listener + httpServer *http.Server + host string + port int + + client *websocket.Conn + clientCtx context.Context + clientCancel context.CancelFunc + writeMu sync.Mutex // serializes writes to client + + // Output coalescing + backpressure + replay (#445). outMu guards all of the + // fields below; outCond signals both the pump (new pending output) and any + // producer blocked by backpressure (pump drained pending). The pump is the + // sole sender of output frames, so replay and live output can never interleave + // or duplicate. + outMu sync.Mutex + outCond *sync.Cond + outPending []byte // coalesced, not-yet-sent output + outClosed bool // set on Stop; releases blocked producers + needReplay bool // a fresh client connected; next flush replays the ring + ring ringBuffer // rolling last-N bytes for late-join replay + pumpCancel context.CancelFunc + pumpDone chan struct{} +} + +// NewServer builds a server. handler may be nil (frames from the client are +// then ignored), which is useful for output-only smoke tests. +func NewServer(cfg Config, handler Handler) *Server { + if cfg.PairTTL <= 0 { + cfg.PairTTL = defaultPairTTL + } + sub, err := fs.Sub(spaFiles, "web") + if err != nil { + // The embed path is a compile-time constant, so this cannot fail in a + // correctly built binary; fall back to the raw FS defensively. + sub = spaFiles + } + s := &Server{ + cfg: cfg, + tokens: NewTokenStore(), + handler: handler, + spa: sub, + state: stateIdle, + ring: ringBuffer{max: replayRingBytes}, + } + s.outCond = sync.NewCond(&s.outMu) + return s +} + +// SetHandler installs the handler that receives client frames. It exists to +// break the construction cycle between the server (a Bridge's Sink) and the +// Bridge (the server's Handler): build the server, build the Bridge with the +// server as Sink, then SetHandler(bridge). It must be called before Start so +// the WebSocket read goroutine never observes a mid-flight swap. +func (s *Server) SetHandler(h Handler) { + s.mu.Lock() + s.handler = h + s.mu.Unlock() +} + +// Start resolves a LAN address, binds a listener, mints a one-time pairing +// token, begins serving, and returns the full pairing URL to print. It is +// non-blocking: the HTTP server runs on its own goroutine. +func (s *Server) Start() (string, error) { + host := s.cfg.Host + if host == "" { + detected, err := DetectRoutableIP() + if err != nil { + return "", err + } + host = detected + } + ln, port, err := ListenFreePort(host, s.cfg.Port) + if err != nil { + return "", err + } + token, err := s.tokens.NewPairing(s.cfg.PairTTL) + if err != nil { + ln.Close() + return "", err + } + + mux := http.NewServeMux() + mux.HandleFunc("/healthz", s.handleHealthz) + mux.HandleFunc("/pair", s.handlePair) + mux.HandleFunc("/ws", s.handleWS) + mux.HandleFunc("/", s.handleRoot) + + pumpCtx, pumpCancel := context.WithCancel(context.Background()) + pumpDone := make(chan struct{}) + + s.mu.Lock() + s.ln = ln + s.host = host + s.port = port + s.httpServer = &http.Server{Handler: mux} + s.state = stateListening + s.pumpCancel = pumpCancel + s.pumpDone = pumpDone + s.mu.Unlock() + + go s.httpServer.Serve(ln) + go s.outputPump(pumpCtx, pumpDone) + + return fmt.Sprintf("http://%s:%d/pair?t=%s", host, port, token), nil +} + +// Stop notifies the client, closes the WebSocket, shuts down the HTTP server, +// and wipes all tokens. It is safe to call more than once. +func (s *Server) Stop(ctx context.Context) error { + s.mu.Lock() + if s.state == stateEnded { + s.mu.Unlock() + return nil + } + s.state = stateEnded + srv := s.httpServer + client := s.client + cancel := s.clientCancel + pumpCancel := s.pumpCancel + pumpDone := s.pumpDone + s.mu.Unlock() + + // Release any producer blocked on backpressure, then stop the output pump and + // wait for its final flush so the last buffered output reaches the client + // before we announce the session end. + s.outMu.Lock() + s.outClosed = true + s.outCond.Broadcast() + s.outMu.Unlock() + if pumpCancel != nil { + pumpCancel() + <-pumpDone + } + + if client != nil { + // Best-effort: tell the browser the session ended, then close. + s.writeFrame(client, Frame{Type: FrameStatus, State: StatusEnded}) + client.Close(websocket.StatusNormalClosure, "session ended") + } + if cancel != nil { + cancel() + } + s.tokens.Clear() + if srv != nil { + return srv.Shutdown(ctx) + } + return nil +} + +func (s *Server) handleHealthz(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) +} + +// handlePair validates the one-time token, issues a session cookie, and +// redirects to the SPA root. Invalid/expired/used tokens get 401. +func (s *Server) handlePair(w http.ResponseWriter, r *http.Request) { + if s.ended() { + http.Error(w, "session ended", http.StatusGone) + return + } + token := r.URL.Query().Get("t") + if token == "" || !s.tokens.ConsumePairing(token) { + http.Error(w, "pairing link invalid or expired", http.StatusUnauthorized) + return + } + cred, err := s.tokens.IssueSession() + if err != nil { + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + http.SetCookie(w, &http.Cookie{ + Name: cookieName, + Value: cred, + Path: "/", + HttpOnly: true, + // Lax, not Strict: the pairing link is opened as a top-level navigation from + // a QR scan (no same-site referrer), and some mobile browsers drop a Strict + // cookie across the /pair→/ redirect that follows — sending the browser to + // the unauthenticated page. Lax is sent on top-level GET navigations (which + // is all this cookie is read on) while still withholding it from cross-site + // subrequests, so it fixes the redirect without weakening the guard. + SameSite: http.SameSiteLaxMode, + }) + http.Redirect(w, r, "/", http.StatusFound) +} + +// handleRoot serves the SPA to authenticated clients and an instructional page +// otherwise. +func (s *Server) handleRoot(w http.ResponseWriter, r *http.Request) { + if s.ended() { + http.Error(w, "session ended", http.StatusGone) + return + } + if !s.authed(r) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`` + + `

Open the pairing link printed in your terminal to connect.

`)) + return + } + http.FileServer(http.FS(s.spa)).ServeHTTP(w, r) +} + +// handleWS upgrades to a WebSocket for the single controlling client. +func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) { + if s.ended() { + http.Error(w, "session ended", http.StatusGone) + return + } + if !s.authed(r) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + s.mu.Lock() + if s.client != nil { + s.mu.Unlock() + http.Error(w, "another device is controlling this session", http.StatusConflict) + return + } + s.mu.Unlock() + + conn, err := websocket.Accept(w, r, nil) + if err != nil { + return + } + conn.SetReadLimit(maxFrameBytes) + + ctx, cancel := context.WithCancel(r.Context()) + s.mu.Lock() + s.client = conn + s.clientCtx = ctx + s.clientCancel = cancel + s.mu.Unlock() + + // Announce connection to the client, then arm a replay so the pump resends the + // rolling scrollback (last replayRingBytes) as the first output frame. Live + // output produced after this point is appended behind the snapshot by the + // pump, so ordering is preserved and nothing is duplicated. + s.writeFrame(conn, Frame{Type: FrameStatus, State: StatusConnected}) + s.outMu.Lock() + s.needReplay = true + s.outCond.Signal() + s.outMu.Unlock() + + // Notify the terminal operator that a client connected (§7.3). + if s.cfg.OnClientConnect != nil { + s.cfg.OnClientConnect(r.RemoteAddr) + } + + defer func() { + cancel() + s.mu.Lock() + if s.client == conn { + s.client = nil + s.clientCtx = nil + s.clientCancel = nil + } + s.mu.Unlock() + conn.Close(websocket.StatusNormalClosure, "") + if s.cfg.OnClientDisconnect != nil { + s.cfg.OnClientDisconnect() + } + }() + + for { + var f Frame + if err := wsjson.Read(ctx, conn, &f); err != nil { + return // client disconnected or context cancelled + } + s.dispatch(f) + } +} + +// dispatch routes an inbound client frame to the handler. +func (s *Server) dispatch(f Frame) { + if s.handler == nil { + return + } + switch f.Type { + case FrameInput: + s.handler.OnInput(f.Text) + case FrameDecide: + s.handler.OnDecide(f.ConfirmID, f.Approve, f.Always) + } +} + +// Broadcast sends a frame to the connected client, if any. It is safe to call +// from any goroutine. +func (s *Server) Broadcast(f Frame) { + s.mu.Lock() + conn := s.client + s.mu.Unlock() + if conn == nil { + return + } + s.writeFrame(conn, f) +} + +// SendOutput streams session text to the client. It never drops bytes: the text +// is appended to the ring (for replay) and to the pending buffer that the pump +// coalesces and flushes. If pending output has grown past outputQueueMax the +// call blocks until the pump drains it, applying backpressure to the producer +// rather than discarding output (§5.4, §8.4). It returns immediately once the +// server is stopping so a shutting-down producer never wedges. +func (s *Server) SendOutput(text string) { + if text == "" { + return + } + b := []byte(text) + + s.outMu.Lock() + defer s.outMu.Unlock() + + // Always record into the ring so a late-joining / reconnecting client can be + // replayed the recent scrollback, even while a producer is momentarily blocked + // on backpressure below. + s.ring.write(b) + + // Backpressure: wait until the pump has drained enough that appending stays + // within the bound. Bail out if the server is stopping. + for !s.outClosed && len(s.outPending) >= outputQueueMax { + s.outCond.Wait() + } + if s.outClosed { + return + } + s.outPending = append(s.outPending, b...) + s.outCond.Signal() +} + +// SendConfirm asks the client to approve a risky tool call. +func (s *Server) SendConfirm(confirmID, tool, summary string) { + s.Broadcast(Frame{Type: FrameConfirm, ConfirmID: confirmID, Tool: tool, Summary: summary}) +} + +// HasClient reports whether a controlling client is currently connected. +func (s *Server) HasClient() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.client != nil +} + +// Addr returns the bound host:port, or "" before Start. +func (s *Server) Addr() string { + s.mu.Lock() + defer s.mu.Unlock() + if s.ln == nil { + return "" + } + return net.JoinHostPort(s.host, fmt.Sprint(s.port)) +} + +func (s *Server) writeFrame(conn *websocket.Conn, f Frame) { + s.writeMu.Lock() + defer s.writeMu.Unlock() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = wsjson.Write(ctx, conn, f) +} + +func (s *Server) ended() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.state == stateEnded +} + +func (s *Server) authed(r *http.Request) bool { + c, err := r.Cookie(cookieName) + if err != nil { + return false + } + return s.tokens.ValidateSession(c.Value) +} + +// outputPump is the sole producer of output frames. It coalesces buffered +// output on a fixed tick and flushes it as a single frame, so a fast stream +// becomes a few batched frames while staying under the latency target: the +// first write after an idle period flushes immediately, and writes arriving +// during the following outputFlushInterval are batched. It also owns replay: +// when a client +// (re)connects, it prepends the ring snapshot before the live pending buffer. +// Being the only writer of output frames, replay and live output can never +// interleave or duplicate. It exits after a final flush when its context is +// cancelled (on Stop). +func (s *Server) outputPump(ctx context.Context, done chan<- struct{}) { + defer close(done) + + ticker := time.NewTicker(outputFlushInterval) + defer ticker.Stop() + + // A tiny goroutine wakes the Cond wait below whenever the context is + // cancelled, so the pump does not sleep past shutdown. + go func() { + <-ctx.Done() + s.outMu.Lock() + s.outCond.Broadcast() + s.outMu.Unlock() + }() + + for { + s.outMu.Lock() + // Wait until there is something to do: pending output, an armed replay, or + // shutdown. + for len(s.outPending) == 0 && !s.needReplay && ctx.Err() == nil { + s.outCond.Wait() + } + replay := s.needReplay + s.needReplay = false + var snapshot []byte + if replay { + snapshot = s.ring.snapshot() + } + // Coalesce: take everything buffered so far as one batch. + pending := s.outPending + s.outPending = nil + // Wake any producer blocked on backpressure now that pending is drained. + s.outCond.Broadcast() + stopping := ctx.Err() != nil + s.outMu.Unlock() + + s.mu.Lock() + conn := s.client + s.mu.Unlock() + + if conn != nil { + // Replay first so the reconnecting browser restores context, then the + // live batch. The ring already contains everything appended via + // SendOutput, so on a fresh connection the snapshot covers pending too; + // avoid double-sending by preferring the snapshot when it is present. + if len(snapshot) > 0 { + s.writeFrame(conn, Frame{Type: FrameOutput, Text: string(snapshot)}) + } else if len(pending) > 0 { + s.writeFrame(conn, Frame{Type: FrameOutput, Text: string(pending)}) + } + } + + if stopping { + // Final drain done; exit. + return + } + + // Rate the loop on the ticker so bursts coalesce instead of spinning. + select { + case <-ctx.Done(): + // Loop once more to perform the final flush of anything buffered while + // we were writing above. + s.finalFlush() + return + case <-ticker.C: + } + } +} + +// finalFlush drains any remaining pending output once during shutdown so the +// last bytes reach the client before the session-ended notice. +func (s *Server) finalFlush() { + s.outMu.Lock() + pending := s.outPending + s.outPending = nil + s.outCond.Broadcast() + s.outMu.Unlock() + if len(pending) == 0 { + return + } + s.mu.Lock() + conn := s.client + s.mu.Unlock() + if conn != nil { + s.writeFrame(conn, Frame{Type: FrameOutput, Text: string(pending)}) + } +} + +// ringBuffer is a bounded rolling byte buffer holding the most recent max bytes +// of session output for late-join / reconnect replay. It is not safe for +// concurrent use; the server guards it with outMu. +type ringBuffer struct { + buf []byte + max int +} + +// write appends p, discarding oldest bytes so the buffer never exceeds max. +func (r *ringBuffer) write(p []byte) { + if r.max <= 0 { + return + } + if len(p) >= r.max { + // Keep only the trailing max bytes of the new data. + r.buf = append(r.buf[:0], p[len(p)-r.max:]...) + return + } + r.buf = append(r.buf, p...) + if len(r.buf) > r.max { + // Drop the oldest overflow. Copy down so the backing array does not grow + // without bound over the life of the session. + drop := len(r.buf) - r.max + r.buf = append(r.buf[:0], r.buf[drop:]...) + } +} + +// snapshot returns a copy of the current contents. +func (r *ringBuffer) snapshot() []byte { + if len(r.buf) == 0 { + return nil + } + out := make([]byte, len(r.buf)) + copy(out, r.buf) + return out +} + +// ErrNotStarted is returned by operations that require a running server. +var ErrNotStarted = errors.New("remotecontrol: server not started") diff --git a/pigo/internal/remotecontrol/server_test.go b/pigo/internal/remotecontrol/server_test.go new file mode 100644 index 0000000..3eeb8f3 --- /dev/null +++ b/pigo/internal/remotecontrol/server_test.go @@ -0,0 +1,464 @@ +package remotecontrol + +import ( + "context" + "net/http" + "net/http/cookiejar" + "strings" + "sync" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" +) + +type fakeHandler struct { + mu sync.Mutex + inputs []string + decides []decision +} + +type decision struct { + id string + approve bool + always bool +} + +func (h *fakeHandler) OnInput(text string) { + h.mu.Lock() + defer h.mu.Unlock() + h.inputs = append(h.inputs, text) +} + +func (h *fakeHandler) OnDecide(id string, approve, always bool) { + h.mu.Lock() + defer h.mu.Unlock() + h.decides = append(h.decides, decision{id, approve, always}) +} + +func (h *fakeHandler) lastInput() string { + h.mu.Lock() + defer h.mu.Unlock() + if len(h.inputs) == 0 { + return "" + } + return h.inputs[len(h.inputs)-1] +} + +// startTestServer boots a server on loopback and returns it plus the pairing +// URL. The caller must Stop it. +func startTestServer(t *testing.T, h Handler) (*Server, string) { + t.Helper() + s := NewServer(Config{Host: "127.0.0.1", Port: 0}, h) + url, err := s.Start() + if err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = s.Stop(ctx) + }) + return s, url +} + +func TestHealthz(t *testing.T) { + _, pairURL := startTestServer(t, nil) + base := pairURL[:strings.Index(pairURL, "/pair")] + resp, err := http.Get(base + "/healthz") + if err != nil { + t.Fatalf("get healthz: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("healthz status = %d, want 200", resp.StatusCode) + } +} + +func TestPairRejectsBadToken(t *testing.T) { + _, pairURL := startTestServer(t, nil) + base := pairURL[:strings.Index(pairURL, "/pair")] + resp, err := http.Get(base + "/pair?t=bogus") + if err != nil { + t.Fatalf("get pair: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("bad-token status = %d, want 401", resp.StatusCode) + } +} + +func TestPairSetsCookieAndServesSPA(t *testing.T) { + _, pairURL := startTestServer(t, nil) + jar, _ := cookiejar.New(nil) + client := &http.Client{Jar: jar} + + resp, err := client.Get(pairURL) + if err != nil { + t.Fatalf("pair: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { // followed redirect to / + t.Fatalf("pair->root status = %d, want 200", resp.StatusCode) + } + + // A second use of the same one-time token must fail. + resp2, err := http.Get(pairURL) + if err != nil { + t.Fatalf("pair reuse: %v", err) + } + defer resp2.Body.Close() + if resp2.StatusCode != http.StatusUnauthorized { + t.Fatalf("token reuse status = %d, want 401", resp2.StatusCode) + } +} + +// The session cookie must be SameSite=Lax so mobile browsers keep it across the +// QR-scan /pair→/ redirect (Strict is dropped by some, breaking pairing). +func TestPairCookieIsSameSiteLax(t *testing.T) { + _, pairURL := startTestServer(t, nil) + client := &http.Client{ + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse // stop at the 302 to read Set-Cookie + }, + } + resp, err := client.Get(pairURL) + if err != nil { + t.Fatalf("pair: %v", err) + } + defer resp.Body.Close() + var got *http.Cookie + for _, c := range resp.Cookies() { + if c.Name == cookieName { + got = c + } + } + if got == nil { + t.Fatal("no session cookie issued") + } + if got.SameSite != http.SameSiteLaxMode { + t.Fatalf("SameSite = %v, want Lax", got.SameSite) + } +} + +func TestRootRequiresAuth(t *testing.T) { + _, pairURL := startTestServer(t, nil) + base := pairURL[:strings.Index(pairURL, "/pair")] + resp, err := http.Get(base + "/") + if err != nil { + t.Fatalf("get root: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("unauth root status = %d, want 401", resp.StatusCode) + } +} + +// sessionCred pairs and extracts the pigo_rc cookie value for WS dialing. +func sessionCred(t *testing.T, pairURL string) (base, cred string) { + t.Helper() + jar, _ := cookiejar.New(nil) + client := &http.Client{ + Jar: jar, + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse // stop at the 302 to read the cookie + }, + } + resp, err := client.Get(pairURL) + if err != nil { + t.Fatalf("pair: %v", err) + } + defer resp.Body.Close() + for _, c := range resp.Cookies() { + if c.Name == cookieName { + cred = c.Value + } + } + if cred == "" { + t.Fatal("no session cookie issued") + } + return pairURL[:strings.Index(pairURL, "/pair")], cred +} + +func dialWS(t *testing.T, base, cred string) (*websocket.Conn, *http.Response, error) { + t.Helper() + wsURL := "ws" + strings.TrimPrefix(base, "http") + "/ws" + opts := &websocket.DialOptions{ + HTTPHeader: http.Header{"Cookie": []string{cookieName + "=" + cred}}, + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + return websocket.Dial(ctx, wsURL, opts) +} + +func TestWSRoundTrip(t *testing.T) { + h := &fakeHandler{} + s, pairURL := startTestServer(t, h) + base, cred := sessionCred(t, pairURL) + + conn, _, err := dialWS(t, base, cred) + if err != nil { + t.Fatalf("dial ws: %v", err) + } + defer conn.Close(websocket.StatusNormalClosure, "") + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + // First frame should be the connected status. + var connected Frame + if err := wsjson.Read(ctx, conn, &connected); err != nil { + t.Fatalf("read connected: %v", err) + } + if connected.Type != FrameStatus || connected.State != StatusConnected { + t.Fatalf("first frame = %+v, want status/connected", connected) + } + + // Client -> server input reaches the handler. + if err := wsjson.Write(ctx, conn, Frame{Type: FrameInput, Text: "hello"}); err != nil { + t.Fatalf("write input: %v", err) + } + deadline := time.Now().Add(time.Second) + for h.lastInput() != "hello" { + if time.Now().After(deadline) { + t.Fatalf("handler never received input, got %q", h.lastInput()) + } + time.Sleep(5 * time.Millisecond) + } + + // Server -> client output reaches the browser. + s.SendOutput("world") + var out Frame + if err := wsjson.Read(ctx, conn, &out); err != nil { + t.Fatalf("read output: %v", err) + } + if out.Type != FrameOutput || out.Text != "world" { + t.Fatalf("output frame = %+v, want output/world", out) + } +} + +func TestWSRejectsUnauth(t *testing.T) { + _, pairURL := startTestServer(t, nil) + base := pairURL[:strings.Index(pairURL, "/pair")] + _, resp, err := dialWS(t, base, "not-a-valid-cred") + if err == nil { + t.Fatal("dial with bad cred succeeded, want failure") + } + if resp != nil && resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401", resp.StatusCode) + } +} + +// readOutput reads frames until it sees a FrameOutput and returns its text, +// skipping any interleaved status frames. +func readOutput(t *testing.T, ctx context.Context, conn *websocket.Conn) string { + t.Helper() + for { + var f Frame + if err := wsjson.Read(ctx, conn, &f); err != nil { + t.Fatalf("read output: %v", err) + } + if f.Type == FrameOutput { + return f.Text + } + } +} + +// TestOutputCoalesced verifies that a burst of writes is coalesced into a +// single output frame by the pump rather than one frame per write, and that no +// bytes are dropped. +func TestOutputCoalesced(t *testing.T) { + s, pairURL := startTestServer(t, &fakeHandler{}) + base, cred := sessionCred(t, pairURL) + + conn, _, err := dialWS(t, base, cred) + if err != nil { + t.Fatalf("dial ws: %v", err) + } + defer conn.Close(websocket.StatusNormalClosure, "") + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + // Drain the connected status frame. + var connected Frame + if err := wsjson.Read(ctx, conn, &connected); err != nil { + t.Fatalf("read connected: %v", err) + } + + // Wait for the server to register the client so writes are not lost before + // the pump has a live connection. + deadline := time.Now().Add(time.Second) + for !s.HasClient() { + if time.Now().After(deadline) { + t.Fatal("client never registered") + } + time.Sleep(2 * time.Millisecond) + } + + // Emit a burst within one flush interval. + const n = 50 + want := "" + for i := 0; i < n; i++ { + s.SendOutput("x") + want += "x" + } + + // Read frames until we have accumulated all the bytes. They must arrive in + // order and total exactly n bytes (no drops, no duplication). Coalescing + // should produce far fewer than n frames. + got := "" + frames := 0 + for len(got) < len(want) { + got += readOutput(t, ctx, conn) + frames++ + } + if got != want { + t.Fatalf("coalesced output = %q, want %q", got, want) + } + if frames >= n { + t.Fatalf("got %d frames for %d writes, expected coalescing", frames, n) + } +} + +// TestReconnectReplay verifies that a client reconnecting mid-session is +// replayed the recent scrollback from the ring buffer. +func TestReconnectReplay(t *testing.T) { + s, pairURL := startTestServer(t, &fakeHandler{}) + base, cred := sessionCred(t, pairURL) + + // First client connects, receives some output, then disconnects. + conn1, _, err := dialWS(t, base, cred) + if err != nil { + t.Fatalf("dial ws 1: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + var connected Frame + if err := wsjson.Read(ctx, conn1, &connected); err != nil { + t.Fatalf("read connected 1: %v", err) + } + deadline := time.Now().Add(time.Second) + for !s.HasClient() { + if time.Now().After(deadline) { + t.Fatal("client 1 never registered") + } + time.Sleep(2 * time.Millisecond) + } + + s.SendOutput("scrollback") + if out := readOutput(t, ctx, conn1); out != "scrollback" { + t.Fatalf("client 1 output = %q, want scrollback", out) + } + conn1.Close(websocket.StatusNormalClosure, "") + + // Wait for the server to release the client slot. + deadline = time.Now().Add(time.Second) + for s.HasClient() { + if time.Now().After(deadline) { + t.Fatal("client 1 slot never released") + } + time.Sleep(2 * time.Millisecond) + } + + // Second client connects and should be replayed the scrollback. + conn2, _, err := dialWS(t, base, cred) + if err != nil { + t.Fatalf("dial ws 2: %v", err) + } + defer conn2.Close(websocket.StatusNormalClosure, "") + if err := wsjson.Read(ctx, conn2, &connected); err != nil { + t.Fatalf("read connected 2: %v", err) + } + if out := readOutput(t, ctx, conn2); out != "scrollback" { + t.Fatalf("replay output = %q, want scrollback", out) + } +} + +// TestClientConnectDisconnectCallbacks verifies the terminal-notice callbacks +// fire on connect and disconnect (§7.3). +func TestClientConnectDisconnectCallbacks(t *testing.T) { + var mu sync.Mutex + var connectedAddr string + connected := make(chan struct{}, 1) + disconnected := make(chan struct{}, 1) + + cfg := Config{ + Host: "127.0.0.1", + Port: 0, + OnClientConnect: func(addr string) { + mu.Lock() + connectedAddr = addr + mu.Unlock() + connected <- struct{}{} + }, + OnClientDisconnect: func() { + disconnected <- struct{}{} + }, + } + s := NewServer(cfg, &fakeHandler{}) + pairURL, err := s.Start() + if err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = s.Stop(ctx) + }) + + base, cred := sessionCred(t, pairURL) + conn, _, err := dialWS(t, base, cred) + if err != nil { + t.Fatalf("dial ws: %v", err) + } + + select { + case <-connected: + case <-time.After(2 * time.Second): + t.Fatal("OnClientConnect never fired") + } + mu.Lock() + addr := connectedAddr + mu.Unlock() + if addr == "" { + t.Fatal("OnClientConnect got empty remote addr") + } + + conn.Close(websocket.StatusNormalClosure, "") + select { + case <-disconnected: + case <-time.After(2 * time.Second): + t.Fatal("OnClientDisconnect never fired") + } +} + +func TestWSSingleClient(t *testing.T) { + s, pairURL := startTestServer(t, &fakeHandler{}) + base, cred := sessionCred(t, pairURL) + + conn1, _, err := dialWS(t, base, cred) + if err != nil { + t.Fatalf("dial first: %v", err) + } + defer conn1.Close(websocket.StatusNormalClosure, "") + + // Wait until the server registers the first client. + deadline := time.Now().Add(time.Second) + for !s.HasClient() { + if time.Now().After(deadline) { + t.Fatal("server never registered first client") + } + time.Sleep(5 * time.Millisecond) + } + + _, resp, err := dialWS(t, base, cred) + if err == nil { + t.Fatal("second client connected, want rejection") + } + if resp != nil && resp.StatusCode != http.StatusConflict { + t.Fatalf("second-client status = %d, want 409", resp.StatusCode) + } +} diff --git a/pigo/internal/remotecontrol/spa_test.go b/pigo/internal/remotecontrol/spa_test.go new file mode 100644 index 0000000..0febe32 --- /dev/null +++ b/pigo/internal/remotecontrol/spa_test.go @@ -0,0 +1,53 @@ +package remotecontrol + +import ( + "io/fs" + "strings" + "testing" +) + +// TestSPAAssetsEmbedded verifies the browser SPA is compiled into the binary +// and exposes the files the server serves. +func TestSPAAssetsEmbedded(t *testing.T) { + sub, err := fs.Sub(spaFiles, "web") + if err != nil { + t.Fatalf("fs.Sub: %v", err) + } + for _, name := range []string{"index.html", "app.js"} { + b, err := fs.ReadFile(sub, name) + if err != nil { + t.Fatalf("embedded %s missing: %v", name, err) + } + if len(b) == 0 { + t.Fatalf("embedded %s is empty", name) + } + } +} + +func TestSPAIndexReferencesApp(t *testing.T) { + b, err := fs.ReadFile(spaFiles, "web/index.html") + if err != nil { + t.Fatalf("read index: %v", err) + } + html := string(b) + for _, want := range []string{`id="output"`, `id="composer"`, `id="confirm"`, "app.js", "viewport"} { + if !strings.Contains(html, want) { + t.Fatalf("index.html missing %q", want) + } + } +} + +func TestSPAScriptHandlesFrames(t *testing.T) { + b, err := fs.ReadFile(spaFiles, "web/app.js") + if err != nil { + t.Fatalf("read app.js: %v", err) + } + js := string(b) + // The client must understand every server->client frame type and emit the + // two client->server types. + for _, want := range []string{`"output"`, `"confirm"`, `"status"`, `type: "input"`, `type: "decide"`, "/ws"} { + if !strings.Contains(js, want) { + t.Fatalf("app.js missing handling for %q", want) + } + } +} diff --git a/pigo/internal/remotecontrol/token.go b/pigo/internal/remotecontrol/token.go new file mode 100644 index 0000000..c7166ae --- /dev/null +++ b/pigo/internal/remotecontrol/token.go @@ -0,0 +1,130 @@ +// Package remotecontrol implements the /remote-control feature: an in-process +// web server that lets a phone on the same LAN mirror the CLI session, inject +// prompts, and approve risky tool calls (see tasks/spec-remote-control.md). +// +// This file (node #438) provides the auth substrate: a one-time, TTL-bound +// pairing token that a paired browser exchanges for an opaque session +// credential. All state is in-memory and process-scoped; nothing is persisted +// or logged, and Clear() wipes it on shutdown. +package remotecontrol + +import ( + "crypto/rand" + "crypto/subtle" + "encoding/hex" + "sync" + "time" +) + +// tokenBytes is the entropy of both pairing tokens and session credentials. +// 32 bytes = 256 bits, hex-encoded to 64 characters in the URL/cookie. +const tokenBytes = 32 + +// pairingToken is a single-use link credential handed to the user (embedded in +// the printed /pair?t=... URL). It is consumed on the first successful pairing +// and rejected thereafter or once expired. +type pairingToken struct { + expiresAt time.Time + used bool +} + +// TokenStore holds the outstanding pairing tokens and issued session +// credentials for one remote-control session. It is safe for concurrent use by +// the HTTP handlers. All values are cryptographically random secrets; the store +// keeps only their hex string form and never logs them. +type TokenStore struct { + mu sync.Mutex + pairing map[string]*pairingToken + sessions map[string]struct{} + // now is the clock, injectable so tests can force expiry without sleeping. + now func() time.Time +} + +// NewTokenStore returns an empty store using the wall clock. +func NewTokenStore() *TokenStore { + return &TokenStore{ + pairing: make(map[string]*pairingToken), + sessions: make(map[string]struct{}), + now: time.Now, + } +} + +// randHex returns n cryptographically random bytes, hex-encoded. crypto/rand +// never returns a short read without an error, so a nil error guarantees a full +// buffer. +func randHex(n int) (string, error) { + buf := make([]byte, n) + if _, err := rand.Read(buf); err != nil { + return "", err + } + return hex.EncodeToString(buf), nil +} + +// NewPairing mints a fresh one-time pairing token that expires after ttl and +// returns its hex value for embedding in the pairing URL. +func (s *TokenStore) NewPairing(ttl time.Duration) (string, error) { + value, err := randHex(tokenBytes) + if err != nil { + return "", err + } + s.mu.Lock() + defer s.mu.Unlock() + s.pairing[value] = &pairingToken{expiresAt: s.now().Add(ttl)} + return value, nil +} + +// ConsumePairing validates a pairing token and, on success, marks it used so it +// can never be redeemed again. It returns false if the token is unknown, +// already used, or expired. The single-use marking happens under the lock, so +// two concurrent /pair requests with the same token cannot both succeed. +func (s *TokenStore) ConsumePairing(value string) bool { + s.mu.Lock() + defer s.mu.Unlock() + tok, ok := s.pairing[value] + if !ok || tok.used || s.now().After(tok.expiresAt) { + return false + } + tok.used = true + return true +} + +// IssueSession creates and stores a new opaque session credential (set as the +// pigo_rc cookie after pairing) and returns its hex value. +func (s *TokenStore) IssueSession() (string, error) { + cred, err := randHex(tokenBytes) + if err != nil { + return "", err + } + s.mu.Lock() + defer s.mu.Unlock() + s.sessions[cred] = struct{}{} + return cred, nil +} + +// ValidateSession reports whether cred matches a currently-issued session +// credential. Comparison is constant-time to avoid leaking the credential +// through response timing; an empty cred is always rejected. +func (s *TokenStore) ValidateSession(cred string) bool { + if cred == "" { + return false + } + want := []byte(cred) + s.mu.Lock() + defer s.mu.Unlock() + var matched bool + for issued := range s.sessions { + if subtle.ConstantTimeCompare([]byte(issued), want) == 1 { + matched = true + } + } + return matched +} + +// Clear wipes all pairing tokens and session credentials. It is called on +// server shutdown so no secret outlives the remote-control session. +func (s *TokenStore) Clear() { + s.mu.Lock() + defer s.mu.Unlock() + s.pairing = make(map[string]*pairingToken) + s.sessions = make(map[string]struct{}) +} diff --git a/pigo/internal/remotecontrol/token_test.go b/pigo/internal/remotecontrol/token_test.go new file mode 100644 index 0000000..4e0714a --- /dev/null +++ b/pigo/internal/remotecontrol/token_test.go @@ -0,0 +1,121 @@ +package remotecontrol + +import ( + "testing" + "time" +) + +func TestNewPairingTokenIsHex256Bit(t *testing.T) { + s := NewTokenStore() + value, err := s.NewPairing(time.Minute) + if err != nil { + t.Fatalf("NewPairing: %v", err) + } + // 32 bytes -> 64 hex chars. + if len(value) != tokenBytes*2 { + t.Fatalf("token length = %d, want %d", len(value), tokenBytes*2) + } + for _, c := range value { + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) { + t.Fatalf("token has non-hex char %q", c) + } + } +} + +func TestPairingIsSingleUse(t *testing.T) { + s := NewTokenStore() + value, err := s.NewPairing(time.Minute) + if err != nil { + t.Fatalf("NewPairing: %v", err) + } + if !s.ConsumePairing(value) { + t.Fatal("first ConsumePairing = false, want true") + } + if s.ConsumePairing(value) { + t.Fatal("second ConsumePairing = true, want false (single-use)") + } +} + +func TestConsumePairingUnknownToken(t *testing.T) { + s := NewTokenStore() + if s.ConsumePairing("deadbeef") { + t.Fatal("ConsumePairing on unknown token = true, want false") + } +} + +func TestPairingExpires(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + s := NewTokenStore() + s.now = func() time.Time { return now } + value, err := s.NewPairing(10 * time.Minute) + if err != nil { + t.Fatalf("NewPairing: %v", err) + } + // Advance clock just past the TTL. + now = now.Add(10*time.Minute + time.Second) + if s.ConsumePairing(value) { + t.Fatal("ConsumePairing after expiry = true, want false") + } +} + +func TestPairingValidJustBeforeExpiry(t *testing.T) { + now := time.Unix(1_700_000_000, 0) + s := NewTokenStore() + s.now = func() time.Time { return now } + value, err := s.NewPairing(10 * time.Minute) + if err != nil { + t.Fatalf("NewPairing: %v", err) + } + now = now.Add(10*time.Minute - time.Second) + if !s.ConsumePairing(value) { + t.Fatal("ConsumePairing just before expiry = false, want true") + } +} + +func TestIssueAndValidateSession(t *testing.T) { + s := NewTokenStore() + cred, err := s.IssueSession() + if err != nil { + t.Fatalf("IssueSession: %v", err) + } + if len(cred) != tokenBytes*2 { + t.Fatalf("cred length = %d, want %d", len(cred), tokenBytes*2) + } + if !s.ValidateSession(cred) { + t.Fatal("ValidateSession(issued) = false, want true") + } + if s.ValidateSession("") { + t.Fatal("ValidateSession(\"\") = true, want false") + } + if s.ValidateSession("not-a-real-cred") { + t.Fatal("ValidateSession(bogus) = true, want false") + } +} + +func TestClearWipesState(t *testing.T) { + s := NewTokenStore() + value, _ := s.NewPairing(time.Minute) + cred, _ := s.IssueSession() + s.Clear() + if s.ConsumePairing(value) { + t.Fatal("pairing token survived Clear") + } + if s.ValidateSession(cred) { + t.Fatal("session credential survived Clear") + } +} + +func TestTokensAreUnique(t *testing.T) { + s := NewTokenStore() + seen := make(map[string]struct{}) + for range 100 { + v, err := s.NewPairing(time.Minute) + if err != nil { + t.Fatalf("NewPairing: %v", err) + } + if _, dup := seen[v]; dup { + t.Fatalf("duplicate token generated: %s", v) + } + seen[v] = struct{}{} + } +} diff --git a/pigo/internal/remotecontrol/web/app.js b/pigo/internal/remotecontrol/web/app.js new file mode 100644 index 0000000..c623202 --- /dev/null +++ b/pigo/internal/remotecontrol/web/app.js @@ -0,0 +1,153 @@ +// pigo remote-control SPA. +// +// Connects to /ws, renders streamed session output, submits prompts, and +// answers confirmation requests. Vanilla JS, no build step — the server embeds +// this file and serves it as a static asset. +(function () { + "use strict"; + + var out = document.getElementById("output"); + var form = document.getElementById("composer"); + var prompt = document.getElementById("prompt"); + var sendBtn = document.getElementById("send"); + var dot = document.getElementById("dot"); + var statusText = document.getElementById("statusText"); + + var confirmEl = document.getElementById("confirm"); + var confirmTool = document.getElementById("confirmTool"); + var confirmSummary = document.getElementById("confirmSummary"); + var confirmAlways = document.getElementById("confirmAlways"); + var btnApprove = document.getElementById("btnApprove"); + var btnReject = document.getElementById("btnReject"); + + var ws = null; + var backoff = 500; // ms, exponential up to 10s + var pendingConfirmId = null; + + // Strip ANSI escape sequences so terminal color codes don't clutter mobile. + var ansi = /\x1b\[[0-9;?]*[ -/]*[@-~]/g; + function clean(s) { return s.replace(ansi, ""); } + + function atBottom() { + return out.scrollHeight - out.scrollTop - out.clientHeight < 40; + } + function appendOutput(text) { + var stick = atBottom(); + out.appendChild(document.createTextNode(clean(text))); + // Cap the scrollback so long sessions don't exhaust memory. + while (out.childNodes.length > 4000) { + out.removeChild(out.firstChild); + } + if (stick) out.scrollTop = out.scrollHeight; + } + + function setStatus(state, label) { + dot.className = state === "on" ? "on" : state === "off" ? "off" : ""; + statusText.textContent = label; + var live = state === "on"; + prompt.disabled = !live; + sendBtn.disabled = !live; + } + + function send(obj) { + if (ws && ws.readyState === WebSocket.OPEN) { + ws.send(JSON.stringify(obj)); + return true; + } + return false; + } + + function showConfirm(frame) { + pendingConfirmId = frame.confirmId; + confirmTool.textContent = frame.tool || "tool"; + confirmSummary.textContent = clean(frame.summary || ""); + confirmAlways.checked = false; + confirmEl.classList.add("show"); + } + function hideConfirm() { + confirmEl.classList.remove("show"); + pendingConfirmId = null; + } + function decide(approve) { + if (pendingConfirmId == null) return; + send({ type: "decide", confirmId: pendingConfirmId, approve: approve, always: confirmAlways.checked }); + hideConfirm(); + } + btnApprove.addEventListener("click", function () { decide(true); }); + btnReject.addEventListener("click", function () { decide(false); }); + + function handleFrame(frame) { + switch (frame.type) { + case "output": + appendOutput(frame.text || ""); + break; + case "confirm": + showConfirm(frame); + break; + case "status": + if (frame.state === "connected") { + setStatus("on", "connected"); + } else if (frame.state === "ended") { + setStatus("off", "session ended"); + } else if (frame.state === "disconnected") { + setStatus("off", frame.reason || "disconnected"); + } + break; + } + } + + function connect() { + var proto = location.protocol === "https:" ? "wss:" : "ws:"; + ws = new WebSocket(proto + "//" + location.host + "/ws"); + + ws.onopen = function () { + backoff = 500; + setStatus("on", "connected"); + }; + ws.onmessage = function (ev) { + var frame; + try { frame = JSON.parse(ev.data); } catch (e) { return; } + handleFrame(frame); + }; + ws.onclose = function () { + setStatus("off", "reconnecting…"); + hideConfirm(); + scheduleReconnect(); + }; + ws.onerror = function () { + try { ws.close(); } catch (e) {} + }; + } + + function scheduleReconnect() { + setTimeout(function () { + backoff = Math.min(backoff * 2, 10000); + connect(); + }, backoff); + } + + // Auto-grow the textarea and submit on Enter (Shift+Enter = newline). + prompt.addEventListener("input", function () { + prompt.style.height = "auto"; + prompt.style.height = Math.min(prompt.scrollHeight, window.innerHeight * 0.4) + "px"; + }); + prompt.addEventListener("keydown", function (e) { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + form.requestSubmit(); + } + }); + + form.addEventListener("submit", function (e) { + e.preventDefault(); + var text = prompt.value; + if (!text.trim()) return; + if (send({ type: "input", text: text })) { + prompt.value = ""; + prompt.style.height = "auto"; + } + }); + + setStatus("", "connecting…"); + connect(); +})(); diff --git a/pigo/internal/remotecontrol/web/index.html b/pigo/internal/remotecontrol/web/index.html new file mode 100644 index 0000000..ffa098a --- /dev/null +++ b/pigo/internal/remotecontrol/web/index.html @@ -0,0 +1,135 @@ + + + + + + + pigo remote + + + +
+ pigo remote + connecting… +
+ +

+
+  
+ + +
+ + + + + + diff --git a/pigo/internal/runtime/args.go b/pigo/internal/runtime/args.go new file mode 100644 index 0000000..f177da5 --- /dev/null +++ b/pigo/internal/runtime/args.go @@ -0,0 +1,34 @@ +// This file implements shell-style argument tokenization for prompt templates +// (US-001, #331). A template invocation like +// +// /component Button "click handler" +// +// arrives at the expander as the raw string `Button "click handler"`; the +// expander needs positional args ($1, $@, ...), so the string must be split into +// ["Button", "click handler"] honoring shell quoting rules: double and single +// quotes group a single argument, surrounding quotes are stripped, and internal +// whitespace is preserved. +// +// Rather than hand-roll a tokenizer, we reuse a mature shell-quoting library +// (github.com/kballard/go-shellquote), per the project's "reuse rather than reinvent" rule. +package runtime + +import "github.com/kballard/go-shellquote" + +// SplitArgs tokenizes a raw argument string using shell-style quoting rules. +// Double quotes ("a b") and single quotes ('a b') each group one argument; +// surrounding quotes are stripped and internal whitespace is preserved. An empty +// (or all-whitespace) input yields an empty non-nil slice with no error. An +// unterminated quote yields an error, so the caller can fall back to treating +// the whole string as $ARGUMENTS rather than feeding a malformed arg list to the +// template engine. +func SplitArgs(s string) ([]string, error) { + parts, err := shellquote.Split(s) + if err != nil { + return nil, err + } + if parts == nil { + return []string{}, nil + } + return parts, nil +} diff --git a/pigo/internal/runtime/args_test.go b/pigo/internal/runtime/args_test.go new file mode 100644 index 0000000..fc03972 --- /dev/null +++ b/pigo/internal/runtime/args_test.go @@ -0,0 +1,69 @@ +package runtime + +// Tests for shell-style argument tokenization (US-001, #331). Covers the cases +// called out in the acceptance criteria: empty input, a single bare argument, +// double-quoted argument with internal space, single-quoted argument, mixed +// quoting, an unterminated quote (error), and leading/trailing whitespace. + +import "testing" + +func TestSplitArgs(t *testing.T) { + cases := []struct { + name string + in string + want []string + }{ + {"empty", "", []string{}}, + {"only whitespace", " ", []string{}}, + {"single bare arg", "Button", []string{"Button"}}, + {"two bare args", "a b", []string{"a", "b"}}, + {"double quoted preserves internal space", `Button "click handler"`, []string{"Button", "click handler"}}, + {"single quoted preserves internal space", `'a b'`, []string{"a b"}}, + {"mixed single and double quotes", `x "y z" 'p q' r`, []string{"x", "y z", "p q", "r"}}, + {"leading and trailing whitespace trimmed", ` hello world `, []string{"hello", "world"}}, + {"quoted arg at boundaries", `"a b" c "d e"`, []string{"a b", "c", "d e"}}, + {"empty double quotes yield empty arg", `""`, []string{""}}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got, err := SplitArgs(c.in) + if err != nil { + t.Fatalf("SplitArgs(%q) returned unexpected error: %v", c.in, err) + } + if len(got) != len(c.want) { + t.Fatalf("SplitArgs(%q) = %v (len %d), want %v (len %d)", c.in, got, len(got), c.want, len(c.want)) + } + for i := range got { + if got[i] != c.want[i] { + t.Errorf("SplitArgs(%q)[%d] = %q, want %q", c.in, i, got[i], c.want[i]) + } + } + }) + } +} + +func TestSplitArgsEmptyReturnsNonNil(t *testing.T) { + got, err := SplitArgs("") + if err != nil { + t.Fatalf("SplitArgs(\"\") returned unexpected error: %v", err) + } + if got == nil { + t.Fatal("SplitArgs(\"\") returned nil slice, want non-nil empty slice") + } + if len(got) != 0 { + t.Fatalf("SplitArgs(\"\") = %v, want empty", got) + } +} + +func TestSplitArgsUnclosedQuoteErrors(t *testing.T) { + cases := []string{ + `"unterminated`, + `'unterminated`, + `foo "bar baz`, + } + for _, in := range cases { + if _, err := SplitArgs(in); err == nil { + t.Errorf("SplitArgs(%q) expected an error for unterminated quote, got nil", in) + } + } +} diff --git a/pigo/internal/runtime/checkpoint.go b/pigo/internal/runtime/checkpoint.go new file mode 100644 index 0000000..4073ba4 --- /dev/null +++ b/pigo/internal/runtime/checkpoint.go @@ -0,0 +1,193 @@ +// Checkpoint persistence for the "infinite context" feature (#480). A checkpoint +// is a distilled summary of a conversation *prefix* — everything up to a +// watermark message index — persisted as a Markdown memory file so a later run +// can reload the collapsed context instead of replaying (and re-tokenizing) the +// whole transcript. +// +// The file lives at /sessions//checkpoint.md and carries +// the repo's standard YAML frontmatter (name/description/metadata.type) with the +// checkpoint bookkeeping (watermark, createdAt, covered message count) under +// metadata; the Summary is the Markdown body. This mirrors the memory-file +// convention (see internal/memory, TypeCheckpoint = "checkpoint") so the memory +// indexer can pick these files up unchanged. +// +// This node provides only the persistence primitives plus the summarize→ +// Checkpoint bridge. Wiring into the run loop (#481) is deliberately out of +// scope: a checkpoint write failure is returned to the caller, which is expected +// to log-and-continue rather than abort the turn. +package runtime + +import ( + "bytes" + "context" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "gopkg.in/yaml.v3" +) + +// Checkpoint is a distilled summary of the conversation up to (and including) +// Watermark. It is the in-memory form written to / read from checkpoint.md. +type Checkpoint struct { + // Watermark is the message index the checkpoint summarizes up to: messages + // [0, Watermark) are collapsed into Summary. A resumed run replays only the + // tail after Watermark, prepending Summary as the collapsed context. + Watermark int + // Summary is the distilled context (the Markdown body of checkpoint.md), + // produced by the summarization LLM call (see compaction.GenerateSummary). + Summary string + // CreatedAt is when the checkpoint was distilled (RFC 3339, UTC). + CreatedAt time.Time + // CoveredMessages is how many messages the summary actually folded in. It + // usually equals Watermark for a linear prefix but is recorded separately so + // a caller that checkpoints a non-contiguous slice still keeps an honest count. + CoveredMessages int +} + +// SummarizeFunc distills a slice of conversation messages into a single summary +// string. It is the seam that lets BuildCheckpoint reuse compaction.GenerateSummary +// without this package depending on the provider stack: the caller supplies a +// closure over GenerateSummary (binding ctx/stream/model/cfg), and BuildCheckpoint +// invokes it. A summarization error is propagated, never swallowed. +type SummarizeFunc func(ctx context.Context, msgs []agentcore.Message) (string, error) + +// checkpointFrontmatter is the YAML head of checkpoint.md. It matches the repo's +// name/description/metadata convention (mirrors SkillFrontmatter and the memory +// file layout) so the file is a well-formed memory document. +type checkpointFrontmatter struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + Metadata checkpointMetadata `yaml:"metadata"` +} + +// checkpointMetadata carries the checkpoint bookkeeping under metadata. Type is +// fixed to "checkpoint" (memory.TypeCheckpoint) so the memory indexer classifies +// it correctly. +type checkpointMetadata struct { + Type string `yaml:"type"` + Watermark int `yaml:"watermark"` + CreatedAt time.Time `yaml:"createdAt"` + CoveredMessages int `yaml:"coveredMessages"` +} + +// checkpointType is the metadata.type value for a checkpoint memory file. It is +// duplicated here (rather than importing internal/memory) to keep this package's +// dependency surface minimal; the two must stay in sync. +const checkpointType = "checkpoint" + +// CheckpointPath returns the on-disk path of a session's checkpoint file: +// /sessions//checkpoint.md. +func CheckpointPath(sessionID, memoryRoot string) string { + return filepath.Join(memoryRoot, "sessions", sessionID, "checkpoint.md") +} + +// BuildCheckpoint distills msgs into a Checkpoint by invoking summarize, tagging +// the result with watermark and now (coerced to UTC). It performs no I/O — the +// caller persists the result with WriteCheckpoint — so the (potentially slow, +// potentially failing) summarization call stays off the write path. A nil +// summarize or a summarization error is returned as an error. +func BuildCheckpoint(ctx context.Context, msgs []agentcore.Message, watermark int, now time.Time, summarize SummarizeFunc) (Checkpoint, error) { + if summarize == nil { + return Checkpoint{}, fmt.Errorf("runtime: BuildCheckpoint: nil summarize func") + } + summary, err := summarize(ctx, msgs) + if err != nil { + return Checkpoint{}, fmt.Errorf("runtime: distill checkpoint: %w", err) + } + return Checkpoint{ + Watermark: watermark, + Summary: summary, + CreatedAt: now.UTC(), + CoveredMessages: len(msgs), + }, nil +} + +// WriteCheckpoint persists cp as /sessions//checkpoint.md, +// creating parent directories (0o755). It writes to a temp file and atomically +// renames it into place so a concurrent reader never sees a half-written file. +// The write is intended to be non-fatal to callers: on error the on-disk file is +// left untouched and the error is returned for the caller to log-and-continue. +func WriteCheckpoint(sessionID, memoryRoot string, cp Checkpoint) error { + if sessionID == "" { + return fmt.Errorf("runtime: WriteCheckpoint: empty sessionID") + } + path := CheckpointPath(sessionID, memoryRoot) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return fmt.Errorf("runtime: create checkpoint dir: %w", err) + } + + doc, err := renderCheckpoint(cp) + if err != nil { + return err + } + + tmp := path + ".tmp" + if err := os.WriteFile(tmp, doc, 0o644); err != nil { + return fmt.Errorf("runtime: write checkpoint temp: %w", err) + } + if err := os.Rename(tmp, path); err != nil { + os.Remove(tmp) + return fmt.Errorf("runtime: commit checkpoint: %w", err) + } + return nil +} + +// renderCheckpoint serializes cp into the checkpoint.md byte form: a "---"-fenced +// YAML frontmatter block followed by the Summary as the Markdown body. +func renderCheckpoint(cp Checkpoint) ([]byte, error) { + fm := checkpointFrontmatter{ + Name: "checkpoint", + Description: fmt.Sprintf("Conversation checkpoint at watermark %d (%d messages).", cp.Watermark, cp.CoveredMessages), + Metadata: checkpointMetadata{ + Type: checkpointType, + Watermark: cp.Watermark, + CreatedAt: cp.CreatedAt.UTC(), + CoveredMessages: cp.CoveredMessages, + }, + } + fmBytes, err := yaml.Marshal(fm) + if err != nil { + return nil, fmt.Errorf("runtime: encode checkpoint frontmatter: %w", err) + } + var b bytes.Buffer + b.WriteString("---\n") + b.Write(fmBytes) + b.WriteString("---\n\n") + b.WriteString(cp.Summary) + return b.Bytes(), nil +} + +// LoadCheckpoint reads and parses the checkpoint for sessionID under memoryRoot. +// It returns (nil, false, nil) when the file does not exist — a missing +// checkpoint is a normal "no collapsed context yet" state, not an error. A +// present-but-malformed file yields a non-nil error. +func LoadCheckpoint(sessionID, memoryRoot string) (*Checkpoint, bool, error) { + path := CheckpointPath(sessionID, memoryRoot) + content, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return nil, false, nil + } + return nil, false, fmt.Errorf("runtime: read checkpoint: %w", err) + } + + fmBytes, body, err := splitFrontmatter(content) + if err != nil { + return nil, false, fmt.Errorf("runtime: parse checkpoint %s: %w", path, err) + } + var fm checkpointFrontmatter + if err := yaml.Unmarshal(fmBytes, &fm); err != nil { + return nil, false, fmt.Errorf("runtime: decode checkpoint frontmatter %s: %w", path, err) + } + + cp := &Checkpoint{ + Watermark: fm.Metadata.Watermark, + Summary: string(bytes.TrimLeft(body, "\r\n")), + CreatedAt: fm.Metadata.CreatedAt.UTC(), + CoveredMessages: fm.Metadata.CoveredMessages, + } + return cp, true, nil +} diff --git a/pigo/internal/runtime/checkpoint_test.go b/pigo/internal/runtime/checkpoint_test.go new file mode 100644 index 0000000..a38a45d --- /dev/null +++ b/pigo/internal/runtime/checkpoint_test.go @@ -0,0 +1,138 @@ +package runtime + +// Tests for checkpoint persistence (#480): the write→load round-trip, the +// missing-file sentinel, and the summarize→Checkpoint bridge. They drive the +// real filesystem via t.TempDir(), matching the session/memory test style. + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// TestCheckpointRoundTrip is the core acceptance check: a written checkpoint +// loads back with watermark, summary, createdAt, and covered count preserved. +func TestCheckpointRoundTrip(t *testing.T) { + root := t.TempDir() + created := time.Date(2026, 8, 1, 9, 30, 0, 0, time.UTC) + cp := Checkpoint{ + Watermark: 42, + Summary: "The user asked to refactor foo().\nWe extracted a helper and added tests.", + CreatedAt: created, + CoveredMessages: 42, + } + if err := WriteCheckpoint("sess-abc", root, cp); err != nil { + t.Fatalf("WriteCheckpoint: %v", err) + } + + got, ok, err := LoadCheckpoint("sess-abc", root) + if err != nil { + t.Fatalf("LoadCheckpoint: %v", err) + } + if !ok { + t.Fatal("LoadCheckpoint: found = false, want true") + } + if got.Watermark != cp.Watermark { + t.Errorf("Watermark = %d, want %d", got.Watermark, cp.Watermark) + } + if got.CoveredMessages != cp.CoveredMessages { + t.Errorf("CoveredMessages = %d, want %d", got.CoveredMessages, cp.CoveredMessages) + } + if got.Summary != cp.Summary { + t.Errorf("Summary = %q, want %q", got.Summary, cp.Summary) + } + if !got.CreatedAt.Equal(cp.CreatedAt) { + t.Errorf("CreatedAt = %v, want %v", got.CreatedAt, cp.CreatedAt) + } +} + +// TestCheckpointFileLayoutAndFrontmatter verifies the file lands at the expected +// path and carries the repo's name/description/metadata.type=checkpoint convention. +func TestCheckpointFileLayoutAndFrontmatter(t *testing.T) { + root := t.TempDir() + cp := Checkpoint{Watermark: 3, Summary: "body text", CreatedAt: time.Now().UTC(), CoveredMessages: 3} + if err := WriteCheckpoint("s1", root, cp); err != nil { + t.Fatalf("WriteCheckpoint: %v", err) + } + want := filepath.Join(root, "sessions", "s1", "checkpoint.md") + if want != CheckpointPath("s1", root) { + t.Fatalf("CheckpointPath = %q, want %q", CheckpointPath("s1", root), want) + } + content, err := os.ReadFile(want) + if err != nil { + t.Fatalf("read checkpoint file: %v", err) + } + fm, _, err := splitFrontmatter(content) + if err != nil { + t.Fatalf("splitFrontmatter: %v", err) + } + s := string(fm) + for _, needle := range []string{"name: checkpoint", "type: checkpoint", "watermark: 3"} { + if !strings.Contains(s, needle) { + t.Errorf("frontmatter missing %q; got:\n%s", needle, s) + } + } +} + +// TestLoadCheckpointMissing verifies a missing checkpoint is the (nil,false,nil) +// sentinel, not an error — callers treat "no checkpoint yet" as normal. +func TestLoadCheckpointMissing(t *testing.T) { + root := t.TempDir() + cp, ok, err := LoadCheckpoint("nope", root) + if err != nil { + t.Fatalf("LoadCheckpoint on missing file: err = %v, want nil", err) + } + if ok { + t.Error("found = true, want false") + } + if cp != nil { + t.Errorf("checkpoint = %+v, want nil", cp) + } +} + +// TestBuildCheckpoint verifies the summarize bridge tags the result with the +// watermark, covered count, and a UTC createdAt. +func TestBuildCheckpoint(t *testing.T) { + msgs := []agentcore.Message{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("hello")}}, + } + now := time.Date(2026, 8, 1, 10, 0, 0, 0, time.FixedZone("x", 3600)) + summarize := func(ctx context.Context, m []agentcore.Message) (string, error) { + if len(m) != len(msgs) { + t.Errorf("summarize got %d msgs, want %d", len(m), len(msgs)) + } + return "distilled", nil + } + cp, err := BuildCheckpoint(context.Background(), msgs, 2, now, summarize) + if err != nil { + t.Fatalf("BuildCheckpoint: %v", err) + } + if cp.Summary != "distilled" { + t.Errorf("Summary = %q, want %q", cp.Summary, "distilled") + } + if cp.Watermark != 2 || cp.CoveredMessages != 2 { + t.Errorf("Watermark/Covered = %d/%d, want 2/2", cp.Watermark, cp.CoveredMessages) + } + if cp.CreatedAt.Location() != time.UTC { + t.Errorf("CreatedAt not UTC: %v", cp.CreatedAt) + } +} + +// TestBuildCheckpointPropagatesError verifies a summarization failure surfaces +// as an error rather than a partial checkpoint (write path never runs). +func TestBuildCheckpointPropagatesError(t *testing.T) { + boom := errors.New("summarize failed") + _, err := BuildCheckpoint(context.Background(), nil, 0, time.Now(), func(context.Context, []agentcore.Message) (string, error) { + return "", boom + }) + if !errors.Is(err, boom) { + t.Fatalf("err = %v, want wrapping %v", err, boom) + } +} diff --git a/pigo/internal/runtime/compaction_test.go b/pigo/internal/runtime/compaction_test.go new file mode 100644 index 0000000..6727e26 --- /dev/null +++ b/pigo/internal/runtime/compaction_test.go @@ -0,0 +1,206 @@ +package runtime + +// Tests for auto-compaction wiring into the loop (US-004, #120): when context +// usage exceeds the usable window after a turn settles, runLoop compacts in +// place and emits a CompactionEvent; a compaction failure is non-fatal and is +// reported via a CompactionEvent carrying an error while the original context is +// preserved. + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/provider" +) + +// bigUserMessages returns n user messages each carrying `chars` characters, used +// to inflate estimated context tokens past a small window. +func bigUserMessages(n, chars int) agentcore.MessageList { + body := strings.Repeat("x", chars) + msgs := make(agentcore.MessageList, 0, n) + for i := 0; i < n; i++ { + msgs = append(msgs, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(body)}, + }) + } + return msgs +} + +// findCompaction returns the first CompactionEvent emitted, or nil. +func findCompaction(events []agentcore.AgentEvent) *agentcore.CompactionEvent { + for _, ev := range events { + if c, ok := ev.(agentcore.CompactionEvent); ok { + return &c + } + } + return nil +} + +// collectEvents drains a stream returning the concrete events (not just kinds). +func collectEvents(t *testing.T, s *LoopEventStream) []agentcore.AgentEvent { + t.Helper() + var out []agentcore.AgentEvent + for ev := range s.Events() { + out = append(out, ev) + } + if _, err := s.Result(context.Background()); err != nil { + t.Fatalf("stream result: %v", err) + } + return out +} + +// summaryStream yields a fixed summary text as an end_turn assistant message, +// standing in for the summarization model. +func summaryStream(text string) provider.StreamFn { + return func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + msg := agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + StopReason: agentcore.StopReasonEndTurn, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + } + s := provider.NewAssistantMessageEventStream(0) + go func() { _ = s.Emit(ctx, provider.StreamDoneEvent{Message: msg}); s.Close() }() + return s, nil + } +} + +func TestAutoCompactionFiresOnThreshold(t *testing.T) { + // Main stream just ends the turn with text; the summary stream is separate so + // the summarization does not consume main-stream scripted turns. + main := scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, + }) + cfg := newRunCfg(main) + cfg.SummaryStream = summaryStream("## Goal\ncompacted") + // Small window + reserve so a handful of fat messages exceed the threshold. + cfg.ContextWindow = 2000 + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + + // Seed a long history so EstimateContextTokens > window-reserve. + agentCtx := &agentcore.AgentContext{Messages: bigUserMessages(12, 800)} + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + ce := findCompaction(events) + if ce == nil { + t.Fatalf("expected a CompactionEvent, got events %+v", events) + } + // A CompactionStartEvent must precede the CompactionEvent so a front-end can + // show an in-progress indicator while summarization is in flight. + var startedBefore bool + for _, ev := range events { + if _, ok := ev.(agentcore.CompactionStartEvent); ok { + startedBefore = true + } + if _, ok := ev.(agentcore.CompactionEvent); ok { + if !startedBefore { + t.Errorf("CompactionStartEvent must be emitted before CompactionEvent") + } + break + } + } + if !startedBefore { + t.Errorf("expected a CompactionStartEvent, got events %+v", events) + } + if ce.ErrorMessage != "" { + t.Fatalf("compaction should have succeeded, got error %q", ce.ErrorMessage) + } + if ce.TokensAfter >= ce.TokensBefore { + t.Errorf("compaction should reduce tokens: before=%d after=%d", ce.TokensBefore, ce.TokensAfter) + } + if ce.SummarizedCount <= 0 { + t.Errorf("expected some messages summarized, got %d", ce.SummarizedCount) + } + // The context must now begin with a compaction checkpoint. + if len(agentCtx.Messages) == 0 || agentCtx.Messages[0].Role() != agentcore.RoleCompaction { + t.Errorf("context should start with a compaction checkpoint, got %+v", agentCtx.Messages) + } +} + +func TestAutoCompactionDisabledWhenWindowUnknown(t *testing.T) { + main := scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, + }) + cfg := newRunCfg(main) + cfg.SummaryStream = summaryStream("unused") + cfg.ContextWindow = 0 // unknown → disabled + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + + agentCtx := &agentcore.AgentContext{Messages: bigUserMessages(12, 800)} + before := len(agentCtx.Messages) + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + if ce := findCompaction(events); ce != nil { + t.Fatalf("no compaction expected when window unknown, got %+v", ce) + } + // Context grows by one assistant reply only (no checkpoint replacement). + if len(agentCtx.Messages) != before+1 { + t.Errorf("context should be untouched by compaction, got %d messages", len(agentCtx.Messages)) + } +} + +func TestCompactionEventEnvelope(t *testing.T) { + env := eventEnvelope(agentcore.CompactionEvent{ + Reason: "threshold", + TokensBefore: 1000, + TokensAfter: 400, + SummarizedCount: 8, + KeptCount: 3, + }) + if env["type"] != agentcore.EventCompaction { + t.Errorf("type = %v, want %q", env["type"], agentcore.EventCompaction) + } + if env["tokensBefore"] != 1000 || env["tokensAfter"] != 400 { + t.Errorf("token fields wrong: %+v", env) + } + if env["summarizedCount"] != 8 || env["keptCount"] != 3 { + t.Errorf("count fields wrong: %+v", env) + } + if _, hasErr := env["error"]; hasErr { + t.Errorf("no error key expected on success: %+v", env) + } + // Failure envelope carries the error. + failEnv := eventEnvelope(agentcore.CompactionEvent{Reason: "threshold", ErrorMessage: "boom"}) + if failEnv["error"] != "boom" { + t.Errorf("error key expected on failure, got %+v", failEnv) + } +} + +func TestAutoCompactionFailureIsNonFatal(t *testing.T) { + main := scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, + }) + cfg := newRunCfg(main) + // Summary stream that fails to build: forces compaction.Compact to error. + cfg.SummaryStream = func(ctx context.Context, model string, llm provider.LlmContext, c provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + return nil, errors.New("summarizer down") + } + cfg.ContextWindow = 2000 + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + + agentCtx := &agentcore.AgentContext{Messages: bigUserMessages(12, 800)} + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + ce := findCompaction(events) + if ce == nil { + t.Fatalf("expected a CompactionEvent reporting the failure") + } + if ce.ErrorMessage == "" { + t.Errorf("failed compaction must carry an ErrorMessage") + } + if ce.TokensAfter != ce.TokensBefore { + t.Errorf("on failure tokens must be unchanged: before=%d after=%d", ce.TokensBefore, ce.TokensAfter) + } + // The run must still end normally. + if events[len(events)-1].EventType() != agentcore.EventAgentEnd { + t.Errorf("run must end with agent_end despite compaction failure") + } + // No compaction checkpoint should have been inserted. + if len(agentCtx.Messages) > 0 && agentCtx.Messages[0].Role() == agentcore.RoleCompaction { + t.Errorf("failed compaction must not insert a checkpoint") + } +} diff --git a/pigo/internal/runtime/config.go b/pigo/internal/runtime/config.go new file mode 100644 index 0000000..7586977 --- /dev/null +++ b/pigo/internal/runtime/config.go @@ -0,0 +1,186 @@ +// This file implements the layered configuration system (US-023, #42), the +// pigo port of pi/zero's config resolution. A resolved Config is produced by +// merging partial layers in precedence order: +// +// default < global < project < environment/CLI +// +// Each layer is a *ConfigLayer whose fields are pointers, so "unset" (nil) is +// distinguishable from "set to the zero value" — only set fields override lower +// layers (field-level replacement, no deep merge). The final Config is +// validated: an unknown thinking level or tool-execution mode is a hard error, +// as is a malformed layer file. +package runtime + +import ( + "encoding/json" + "fmt" + "os" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" +) + +// Config is the fully resolved configuration a run operates under, after all +// layers are merged and validated. +type Config struct { + // Model is the default model id used when a run does not specify one. + Model string + // Provider is the default provider name. + Provider string + // Credentials maps provider name → API key. Merged per-provider across + // layers (a higher layer's key for provider X overrides a lower one, but + // providers only present in a lower layer are retained). + Credentials map[string]string + // ToolExecutionMode is the default execution mode for tools that do not pin + // their own mode. + ToolExecutionMode agentcore.ToolExecutionMode + // ThinkingLevel is the default reasoning-effort level. + ThinkingLevel agentcore.ThinkingLevel + // Hooks are the resolved, append-merged hook matchers keyed by event type. + // Nil when no layer defined any hooks (FR-18), so the no-hooks path costs + // nothing downstream. + Hooks hooks.HookSet +} + +// ConfigLayer is one partial layer of configuration. Pointer/optional fields +// distinguish "not set in this layer" (nil/empty) from an explicit value, so a +// higher layer only overrides the fields it actually sets. +type ConfigLayer struct { + Model *string `json:"model,omitempty"` + Provider *string `json:"provider,omitempty"` + Credentials map[string]string `json:"credentials,omitempty"` + ToolExecutionMode *string `json:"toolExecutionMode,omitempty"` + ThinkingLevel *string `json:"thinkingLevel,omitempty"` + // Hooks are this layer's hook matchers keyed by event type. Unlike the + // scalar fields, hooks are not overridden across layers: ResolveConfig + // appends each layer's matchers per event type (FR-2), so lower-layer hooks + // always still fire. + Hooks hooks.HookSet `json:"hooks,omitempty"` +} + +// DefaultConfigLayer is the base layer applied before all others. It gives a +// usable configuration out of the box. +func DefaultConfigLayer() ConfigLayer { + model := "openrouter/free" + provider := "openrouter" + mode := string(agentcore.ToolExecutionParallel) + level := string(agentcore.ThinkingMedium) + return ConfigLayer{ + Model: &model, + Provider: &provider, + ToolExecutionMode: &mode, + ThinkingLevel: &level, + } +} + +// LoadConfigLayer reads and decodes a single JSON config layer from path. A +// missing file yields a nil layer and no error (an absent layer is not a +// failure); a present-but-malformed file is a hard error. +func LoadConfigLayer(path string) (*ConfigLayer, error) { + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, fmt.Errorf("read config %s: %w", path, err) + } + var layer ConfigLayer + if err := json.Unmarshal(data, &layer); err != nil { + return nil, fmt.Errorf("parse config %s: %w", path, err) + } + return &layer, nil +} + +// ResolveConfig merges the given layers in ascending precedence order (earlier +// layers are overridden by later ones) and validates the result. Nil layers are +// skipped, so callers can pass the output of LoadConfigLayer directly. The +// merge is field-level: a later layer only overrides fields it sets, except +// Credentials, which merges per-provider. +func ResolveConfig(layers ...*ConfigLayer) (Config, error) { + var cfg Config + for _, layer := range layers { + if layer == nil { + continue + } + if layer.Model != nil { + cfg.Model = *layer.Model + } + if layer.Provider != nil { + cfg.Provider = *layer.Provider + } + if layer.ToolExecutionMode != nil { + cfg.ToolExecutionMode = agentcore.ToolExecutionMode(*layer.ToolExecutionMode) + } + if layer.ThinkingLevel != nil { + cfg.ThinkingLevel = agentcore.ThinkingLevel(*layer.ThinkingLevel) + } + for provider, key := range layer.Credentials { + if cfg.Credentials == nil { + cfg.Credentials = make(map[string]string) + } + cfg.Credentials[provider] = key + } + for eventType, matchers := range layer.Hooks { + if len(matchers) == 0 { + continue + } + if cfg.Hooks == nil { + cfg.Hooks = make(hooks.HookSet) + } + cfg.Hooks[eventType] = append(cfg.Hooks[eventType], matchers...) + } + } + if err := cfg.validate(); err != nil { + return Config{}, err + } + return cfg, nil +} + +// EnvConfigLayer builds a config layer from environment variables, the highest +// file-independent layer (below only explicit CLI flags). Recognized: +// +// PIGO_MODEL, PIGO_PROVIDER, PIGO_TOOL_EXECUTION_MODE, PIGO_THINKING_LEVEL +// +// Only set variables contribute; unset ones leave the field nil so lower layers +// show through. Credential env vars are intentionally NOT captured here — keys +// are resolved lazily by the CredentialStore and never merged into a struct +// that might be logged (US-012). +func EnvConfigLayer(getenv func(string) string) ConfigLayer { + if getenv == nil { + getenv = os.Getenv + } + var layer ConfigLayer + if v := getenv("PIGO_MODEL"); v != "" { + layer.Model = &v + } + if v := getenv("PIGO_PROVIDER"); v != "" { + layer.Provider = &v + } + if v := getenv("PIGO_TOOL_EXECUTION_MODE"); v != "" { + layer.ToolExecutionMode = &v + } + if v := getenv("PIGO_THINKING_LEVEL"); v != "" { + layer.ThinkingLevel = &v + } + return layer +} + +// validate reports the first invalid field in the resolved config: an unknown +// tool-execution mode or thinking level. An empty model is also rejected, since +// a run cannot proceed without one. +func (c Config) validate() error { + if c.Model == "" { + return fmt.Errorf("config: model must not be empty") + } + switch c.ToolExecutionMode { + case agentcore.ToolExecutionParallel, agentcore.ToolExecutionSequential: + default: + return fmt.Errorf("config: invalid toolExecutionMode %q (want parallel|sequential)", c.ToolExecutionMode) + } + switch c.ThinkingLevel { + case agentcore.ThinkingOff, agentcore.ThinkingMinimal, agentcore.ThinkingLow, agentcore.ThinkingMedium, agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax: + default: + return fmt.Errorf("config: invalid thinkingLevel %q", c.ThinkingLevel) + } + return nil +} diff --git a/pigo/internal/runtime/config_test.go b/pigo/internal/runtime/config_test.go new file mode 100644 index 0000000..8765819 --- /dev/null +++ b/pigo/internal/runtime/config_test.go @@ -0,0 +1,240 @@ +package runtime + +// Tests for the layered configuration system (US-023, #42): the precedence +// order (default < global < project < env/CLI), per-provider credential merge, +// and the hard-error paths for malformed files and invalid field values. + +import ( + "os" + "path/filepath" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/hooks" +) + +// ptr is a helper for building pointer-valued config-layer fields in tests. +func ptr[T any](v T) *T { return &v } + +// TestResolveConfigPrecedence is the acceptance-critical test: a field set in a +// higher layer overrides the same field in every lower layer, and the winner is +// always the highest layer that set it. +func TestResolveConfigPrecedence(t *testing.T) { + def := DefaultConfigLayer() + global := &ConfigLayer{Model: ptr("global/model"), ThinkingLevel: ptr("low")} + project := &ConfigLayer{Model: ptr("project/model")} + env := &ConfigLayer{Model: ptr("env/model"), Provider: ptr("bedrock")} + + cfg, err := ResolveConfig(&def, global, project, env) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + // Model set in all four layers → env (highest) wins. + if cfg.Model != "env/model" { + t.Errorf("Model = %q, want env/model (highest layer wins)", cfg.Model) + } + // Provider set only in env → env value. + if cfg.Provider != "bedrock" { + t.Errorf("Provider = %q, want bedrock", cfg.Provider) + } + // ThinkingLevel set in global only (not project/env) → global value shows through. + if cfg.ThinkingLevel != agentcore.ThinkingLow { + t.Errorf("ThinkingLevel = %q, want low (from global, lower layers don't set it)", cfg.ThinkingLevel) + } + // ToolExecutionMode set only in default → default shows through. + if cfg.ToolExecutionMode != agentcore.ToolExecutionParallel { + t.Errorf("ToolExecutionMode = %q, want parallel (default)", cfg.ToolExecutionMode) + } +} + +// TestResolveConfigCredentialMerge verifies credentials merge per-provider: a +// higher layer overrides one provider's key while a lower layer's other-provider +// key is retained. +func TestResolveConfigCredentialMerge(t *testing.T) { + def := DefaultConfigLayer() + global := &ConfigLayer{Credentials: map[string]string{"openrouter": "or-low", "ollama": "ol-key"}} + project := &ConfigLayer{Credentials: map[string]string{"openrouter": "or-high"}} + + cfg, err := ResolveConfig(&def, global, project) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + if cfg.Credentials["openrouter"] != "or-high" { + t.Errorf("openrouter key = %q, want or-high (project overrides global)", cfg.Credentials["openrouter"]) + } + if cfg.Credentials["ollama"] != "ol-key" { + t.Errorf("ollama key = %q, want ol-key (retained from global)", cfg.Credentials["ollama"]) + } +} + +// TestResolveConfigInvalidValues verifies invalid field values are hard errors. +func TestResolveConfigInvalidValues(t *testing.T) { + def := DefaultConfigLayer() + cases := []struct { + name string + layer *ConfigLayer + }{ + {"bad mode", &ConfigLayer{ToolExecutionMode: ptr("concurrent")}}, + {"bad thinking", &ConfigLayer{ThinkingLevel: ptr("ultra")}}, + {"empty model", &ConfigLayer{Model: ptr("")}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if _, err := ResolveConfig(&def, tc.layer); err == nil { + t.Errorf("%s must be a hard error, got nil", tc.name) + } + }) + } +} + +// TestLoadConfigLayerMissingAndMalformed verifies a missing file is not an error +// (nil layer) while a malformed file is. +func TestLoadConfigLayerMissingAndMalformed(t *testing.T) { + dir := t.TempDir() + + // Missing file → nil, nil. + layer, err := LoadConfigLayer(filepath.Join(dir, "nope.json")) + if err != nil || layer != nil { + t.Errorf("missing file: got (%v, %v), want (nil, nil)", layer, err) + } + + // Malformed file → error. + bad := filepath.Join(dir, "bad.json") + if err := os.WriteFile(bad, []byte("{not json"), 0o600); err != nil { + t.Fatal(err) + } + if _, err := LoadConfigLayer(bad); err == nil { + t.Error("malformed config file must return an error") + } + + // Well-formed file → decoded layer. + good := filepath.Join(dir, "good.json") + if err := os.WriteFile(good, []byte(`{"model":"m","provider":"p"}`), 0o600); err != nil { + t.Fatal(err) + } + layer, err = LoadConfigLayer(good) + if err != nil { + t.Fatalf("good file: %v", err) + } + if layer == nil || layer.Model == nil || *layer.Model != "m" { + t.Errorf("good file decoded incorrectly: %+v", layer) + } +} + +// TestEnvConfigLayer verifies env vars map to the right fields and unset vars +// leave fields nil. +func TestEnvConfigLayer(t *testing.T) { + env := map[string]string{ + "PIGO_MODEL": "env/m", + "PIGO_THINKING_LEVEL": "high", + "PIGO_TOOL_EXECUTION_MODE": "sequential", + } + layer := EnvConfigLayer(func(k string) string { return env[k] }) + if layer.Model == nil || *layer.Model != "env/m" { + t.Errorf("PIGO_MODEL not captured: %+v", layer.Model) + } + if layer.ThinkingLevel == nil || *layer.ThinkingLevel != "high" { + t.Errorf("PIGO_THINKING_LEVEL not captured: %+v", layer.ThinkingLevel) + } + if layer.ToolExecutionMode == nil || *layer.ToolExecutionMode != "sequential" { + t.Errorf("PIGO_TOOL_EXECUTION_MODE not captured: %+v", layer.ToolExecutionMode) + } + // Unset var → nil field. + if layer.Provider != nil { + t.Errorf("unset PIGO_PROVIDER should leave Provider nil, got %v", *layer.Provider) + } +} + +// TestResolveConfigDefaultsAlone verifies the default layer alone yields a valid +// config. +func TestResolveConfigDefaultsAlone(t *testing.T) { + def := DefaultConfigLayer() + cfg, err := ResolveConfig(&def) + if err != nil { + t.Fatalf("default-only config must be valid: %v", err) + } + if cfg.Model == "" || cfg.ToolExecutionMode == "" || cfg.ThinkingLevel == "" { + t.Errorf("default config incomplete: %+v", cfg) + } +} + +// TestResolveConfigHooksAppendMerge verifies hooks are append-merged per event +// type across layers (FR-2) in ascending order, not overridden like scalars. +func TestResolveConfigHooksAppendMerge(t *testing.T) { + def := DefaultConfigLayer() + global := &ConfigLayer{Hooks: hooks.HookSet{ + "PreToolUse": {{Matcher: "*", Hooks: []hooks.HookConfig{{Command: "global-pre"}}}}, + "Stop": {{Matcher: "", Hooks: []hooks.HookConfig{{Command: "global-stop"}}}}, + }} + project := &ConfigLayer{Hooks: hooks.HookSet{ + "PreToolUse": {{Matcher: "bash", Hooks: []hooks.HookConfig{{Command: "project-pre"}}}}, + }} + + cfg, err := ResolveConfig(&def, global, project) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + pre := cfg.Hooks["PreToolUse"] + if len(pre) != 2 { + t.Fatalf("PreToolUse matchers = %d, want 2 (global+project appended)", len(pre)) + } + // Ascending layer order: global before project. + if pre[0].Hooks[0].Command != "global-pre" || pre[1].Hooks[0].Command != "project-pre" { + t.Errorf("PreToolUse order wrong: %q, %q", pre[0].Hooks[0].Command, pre[1].Hooks[0].Command) + } + if len(cfg.Hooks["Stop"]) != 1 { + t.Errorf("Stop matchers = %d, want 1 (only global set it)", len(cfg.Hooks["Stop"])) + } +} + +// TestResolveConfigHooksNilWhenAbsent verifies the no-hooks path leaves +// cfg.Hooks nil (FR-18), so downstream can cheaply skip hook dispatch. +func TestResolveConfigHooksNilWhenAbsent(t *testing.T) { + def := DefaultConfigLayer() + cfg, err := ResolveConfig(&def) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + if cfg.Hooks != nil { + t.Errorf("Hooks = %v, want nil when no layer defines hooks", cfg.Hooks) + } + // An empty (but non-nil) hook map in a layer must not allocate cfg.Hooks. + empty := &ConfigLayer{Hooks: hooks.HookSet{"PreToolUse": {}}} + cfg2, err := ResolveConfig(&def, empty) + if err != nil { + t.Fatalf("ResolveConfig: %v", err) + } + if cfg2.Hooks != nil { + t.Errorf("Hooks = %v, want nil when layer's event has no matchers", cfg2.Hooks) + } +} + +// TestLoadConfigLayerHooks verifies a layer's hooks decode from JSON, including +// per-hook timeout and matcher fields. +func TestLoadConfigLayerHooks(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.json") + body := `{ + "model": "x/y", + "hooks": { + "PreToolUse": [ + {"matcher": "bash", "hooks": [{"type": "command", "command": "echo hi", "timeout": 5}]} + ] + } + }` + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + layer, err := LoadConfigLayer(path) + if err != nil { + t.Fatalf("LoadConfigLayer: %v", err) + } + pre := layer.Hooks["PreToolUse"] + if len(pre) != 1 || pre[0].Matcher != "bash" { + t.Fatalf("unexpected matchers: %+v", pre) + } + h := pre[0].Hooks[0] + if h.Command != "echo hi" || h.TimeoutSeconds() != 5 { + t.Errorf("unexpected hook: cmd=%q timeout=%d", h.Command, h.TimeoutSeconds()) + } +} diff --git a/pigo/internal/runtime/e2e_robustness_test.go b/pigo/internal/runtime/e2e_robustness_test.go new file mode 100644 index 0000000..d2f0e9b --- /dev/null +++ b/pigo/internal/runtime/e2e_robustness_test.go @@ -0,0 +1,149 @@ +package runtime + +// End-to-end robustness verification scenarios (US-007 / FR-10): the harness's +// two self-protection mechanisms must fire correctly when driven through the +// real provider seam with NO external LLM. +// +// - LONG SESSION: a run whose accumulated context exceeds the usable window +// (ContextWindow − ReserveTokens) must auto-compact in place, emitting a +// successful CompactionEvent, and still finish normally. +// - LARGE OUTPUT: a tool that returns far more than the executor-layer byte +// budget (toolResultMaxBytes = 100_000) must have its result truncated with +// the shared "[truncated" annotation, so a single fat tool result cannot +// overflow the model context, and the run still finishes. +// +// Both scenarios reuse the existing in-repo test infrastructure only: the faux +// provider seam (StreamFnFromProvider via newFauxRunCfg / toolCallTurn / textTurn), +// the scripted StreamFn (scriptedStream / newRunCfg / summaryStream), and the +// event collectors (collectEvents / collectStream / findCompaction). No new +// mocking is invented and no production (non-_test) code is touched. + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/compaction" +) + +// TestE2E_LongSession_TriggersCompaction drives a long session whose seeded +// history already exceeds the usable window, forcing the loop to auto-compact +// after the first turn settles. It asserts a *successful* CompactionEvent is +// emitted (empty ErrorMessage), tokens shrink, a compaction checkpoint replaces +// the head of the context, and the run ends normally via agent_end. +func TestE2E_LongSession_TriggersCompaction(t *testing.T) { + // Main stream just ends the turn; a separate summary stream stands in for the + // summarization model so compaction does not consume main-stream turns. + main := scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("ack")}}, + }) + cfg := newRunCfg(main) + cfg.SummaryStream = summaryStream("## Goal\nlong session compacted") + // Deliberately tiny window so a handful of fat seeded messages exceed the + // threshold (ContextWindow − ReserveTokens) the moment the first turn settles. + cfg.ContextWindow = 2000 + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + + // Seed a long history: 16 messages × 800 chars each blows past the usable window. + agentCtx := &agentcore.AgentContext{Messages: bigUserMessages(16, 800)} + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + + ce := findCompaction(events) + if ce == nil { + t.Fatalf("long session must trigger a CompactionEvent, got %v", eventKinds(events)) + } + if ce.ErrorMessage != "" { + t.Fatalf("compaction must succeed, got error %q", ce.ErrorMessage) + } + if ce.TokensAfter >= ce.TokensBefore { + t.Errorf("compaction must reduce estimated tokens: before=%d after=%d", ce.TokensBefore, ce.TokensAfter) + } + if ce.SummarizedCount <= 0 { + t.Errorf("compaction must fold at least one message into the summary, got %d", ce.SummarizedCount) + } + // The compacted context must begin with a compaction checkpoint. + if len(agentCtx.Messages) == 0 || agentCtx.Messages[0].Role() != agentcore.RoleCompaction { + t.Errorf("context must start with a compaction checkpoint after compaction, got %+v", agentCtx.Messages) + } + // The run must still terminate cleanly. + if n := len(events); n == 0 || events[n-1].EventType() != agentcore.EventAgentEnd { + t.Errorf("run must end with agent_end, got %v", eventKinds(events)) + } +} + +// TestE2E_LargeOutput_TriggersTruncation drives a tool call whose tool returns +// output far larger than the executor-layer byte budget. It asserts the +// resulting tool-result content carries the shared "[truncated" annotation and +// is clipped well below the raw size (so the context cannot overflow from a +// single fat result), and the run still completes normally. +func TestE2E_LargeOutput_TriggersTruncation(t *testing.T) { + // A payload well over toolResultMaxBytes (100_000). 250_000 bytes guarantees + // the executor-layer budget bites regardless of any looser inner cap. + const rawSize = 250_000 + huge := strings.Repeat("A", rawSize) + + // A tool that emits the oversized payload as a single text block. + bigOutputTool := execTool{ + name: "flood", + mode: agentcore.ToolExecutionParallel, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(huge)}}, nil + }, + } + + // Turn 1: call the flooding tool. Turn 2: end the turn with text. + p := &fauxProvider{ + name: "faux", + turns: []fauxTurn{ + toolCallTurn("call-flood", "flood", `{}`), + textTurn("handled large output"), + }, + } + cfg := newFauxRunCfg(p, bigOutputTool) + // A generous window: the point is that truncation keeps the result small + // enough that the context does NOT overflow, so no compaction is needed. + cfg.ContextWindow = 200_000 + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("run flood")}}}} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + // Locate the flood tool result and assert it was truncated. + var floodResult *agentcore.ToolResultMessage + for i := range msgs { + if tr, ok := msgs[i].(agentcore.ToolResultMessage); ok && tr.ToolCallID == "call-flood" { + trCopy := tr + floodResult = &trCopy + } + } + if floodResult == nil { + t.Fatalf("expected a tool result for the flood call, got %+v", msgs) + } + got := textContentOf(floodResult.Content) + if !strings.Contains(got, "[truncated") { + t.Errorf("large tool output must carry the \"[truncated\" annotation, got %d bytes without it", len(got)) + } + // The clipped result must be far smaller than the raw payload — the context + // protection actually reduced the size (it must not blow the 100_000 budget + // wildly; allow generous headroom for head+tail+marker). + if len(got) >= rawSize { + t.Errorf("truncation must shrink the result: got %d bytes, raw was %d", len(got), rawSize) + } + if len(got) > 120_000 { + t.Errorf("truncated result should be near the byte budget, got %d bytes", len(got)) + } + + // Context must not have overflowed: with truncation in place the tiny clipped + // result never crosses the window, so no compaction should have fired. + for _, ev := range kinds { + if ev == agentcore.EventCompaction { + t.Errorf("truncation should keep context under the window; no compaction expected, got kinds %v", kinds) + } + } + // And the run finished cleanly. + if len(kinds) == 0 || kinds[len(kinds)-1] != agentcore.EventAgentEnd { + t.Errorf("run must end with agent_end, got %v", kinds) + } +} diff --git a/pigo/internal/runtime/faux_provider_test.go b/pigo/internal/runtime/faux_provider_test.go new file mode 100644 index 0000000..bf176f3 --- /dev/null +++ b/pigo/internal/runtime/faux_provider_test.go @@ -0,0 +1,429 @@ +package runtime + +// This file implements the faux provider (mirrors pi providers/faux.ts) and the +// loop integration tests that drive the whole agent loop through it — the +// project's primary and only core test seam (US-002 / Testing Decisions, #16). +// +// Unlike loop_test.go, which drives the loop with a coarse StreamFn that emits +// only a terminal StreamDoneEvent, the faux provider is a real Provider whose +// StreamCompletion replays a *fine-grained* script of AssistantMessageEvents +// (start → text/toolcall deltas → done) — one scripted turn per call. It is +// wired into the loop via StreamFnFromProvider, the real seam, so the whole +// path (message_start / message_update / message_end deltas, the six hooks, +// truncation protection, parallel ordering, and EventStream cancellation) is +// covered end to end without mocking any loop-internal function. + +import ( + "context" + "encoding/json" + "sync" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/provider" +) + +// fauxTurn is one scripted assistant turn: the fine-grained stream events the +// faux provider replays for a single StreamCompletion call. +type fauxTurn []provider.AssistantMessageEvent + +// textTurn scripts a turn that streams text as start → text delta → done(end_turn). +func textTurn(text string) fauxTurn { + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + withText := partial + withText.Content = agentcore.ContentList{agentcore.NewTextContent(text)} + final := withText + final.StopReason = agentcore.StopReasonEndTurn + return fauxTurn{ + provider.StreamStartEvent{Partial: partial}, + provider.StreamTextEvent{Partial: withText}, + provider.StreamDoneEvent{Message: final}, + } +} + +// toolCallTurn scripts a turn that streams one tool call as +// start → toolcall delta → done(tool_use). +func toolCallTurn(id, name, args string) fauxTurn { + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + withCall := partial + withCall.Content = agentcore.ContentList{agentcore.NewToolCallContent(id, name, json.RawMessage(args))} + final := withCall + final.StopReason = agentcore.StopReasonToolUse + return fauxTurn{ + provider.StreamStartEvent{Partial: partial}, + provider.StreamToolCallEvent{Partial: withCall}, + provider.StreamDoneEvent{Message: final}, + } +} + +// fauxProvider is a real Provider that replays one scripted turn per +// StreamCompletion call, in order. It records every request it received so +// tests can assert what the loop actually sent (model, context, config). Once +// the script is exhausted it replays a plain end_turn turn. +type fauxProvider struct { + name string + models []provider.Model + turns []fauxTurn + + mu sync.Mutex + calls int + requests []provider.CompletionRequest + // delay optionally slows each delta emit, used by the cancellation test to + // keep the stream open long enough to cancel mid-flight. + delay time.Duration +} + +func (p *fauxProvider) Name() string { return p.name } +func (p *fauxProvider) Models() []provider.Model { return p.models } + +func (p *fauxProvider) callCount() int { + p.mu.Lock() + defer p.mu.Unlock() + return p.calls +} + +func (p *fauxProvider) requestAt(i int) provider.CompletionRequest { + p.mu.Lock() + defer p.mu.Unlock() + return p.requests[i] +} + +func (p *fauxProvider) StreamCompletion(ctx context.Context, req provider.CompletionRequest) (*provider.AssistantMessageEventStream, error) { + p.mu.Lock() + idx := p.calls + p.calls++ + p.requests = append(p.requests, req) + var turn fauxTurn + if idx < len(p.turns) { + turn = p.turns[idx] + } else { + turn = textTurn("") + } + delay := p.delay + p.mu.Unlock() + + s := provider.NewAssistantMessageEventStream(0) + go func() { + for _, ev := range turn { + if delay > 0 { + select { + case <-time.After(delay): + case <-ctx.Done(): + s.SetError(ctx.Err()) + s.Close() + return + } + } + if err := s.Emit(ctx, ev); err != nil { + s.SetError(err) + s.Close() + return + } + } + s.Close() + }() + return s, nil +} + +// newFauxRunCfg wires a faux provider into the loop via StreamFnFromProvider +// (the real seam) and registers the given tools. No loop-internal function is +// mocked — only the provider boundary. +func newFauxRunCfg(p *fauxProvider, tools ...agentcore.AgentTool) RunConfig { + reg := agenttool.NewToolRegistry() + for _, tl := range tools { + _ = reg.Register(tl) + } + return RunConfig{ + LoopConfig: LoopConfig{Model: "faux", Stream: provider.StreamFnFromProvider(p)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}}, + } +} + +// TestFauxProviderTextToolText drives the flagship seam scenario end to end: +// text → tool call → text over the real loop, asserting both the AgentEvent +// stream shape and the final []AgentMessage. Nothing loop-internal is mocked. +func TestFauxProviderTextToolText(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{ + textTurn("thinking about it"), // turn 1: plain text, no tool + toolCallTurn("call-1", "echo", `{"msg":"hello"}`), // turn 2: tool call + textTurn("all done"), // turn 3: final text + }, + } + // GetFollowUpMessages injects a follow-up once so the loop advances past the + // first natural (text-only) turn end into the tool-call turn. + served := false + cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false)) + cfg.GetFollowUpMessages = func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage { + if served { + return nil + } + served = true + return []agentcore.AgentMessage{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("go on")}}} + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + // Event shape: message deltas must appear (start/update/end), a tool + // executed exactly once, and the run bookended by agent_start/agent_end. + if kinds[0] != agentcore.EventAgentStart || kinds[len(kinds)-1] != agentcore.EventAgentEnd { + t.Fatalf("run must be bracketed by agent_start/agent_end, got %v", kinds) + } + if countKind(kinds, agentcore.EventMessageStart) < 3 || countKind(kinds, agentcore.EventMessageEnd) < 3 { + t.Errorf("expected fine-grained message deltas for each turn, got %v", kinds) + } + if countKind(kinds, agentcore.EventMessageUpdate) < 3 { + t.Errorf("expected message_update deltas (text/toolcall), got %v", kinds) + } + if got := countKind(kinds, agentcore.EventToolExecutionStart); got != 1 { + t.Errorf("expected 1 tool_execution_start, got %d in %v", got, kinds) + } + if got := countKind(kinds, agentcore.EventToolExecutionEnd); got != 1 { + t.Errorf("expected 1 tool_execution_end, got %d in %v", got, kinds) + } + if got := countKind(kinds, agentcore.EventTurnStart); got != 3 { + t.Errorf("expected 3 turns (text→tool→text), got %d in %v", got, kinds) + } + + // Final messages: assistant(text) + user(follow-up) + assistant(tool) + + // toolResult + assistant(text) = 5, in order. + if len(msgs) != 5 { + t.Fatalf("expected 5 new messages, got %d: %+v", len(msgs), msgs) + } + if a, ok := msgs[0].(agentcore.AssistantMessage); !ok || textContentOf(a.Content) != "thinking about it" { + t.Errorf("msg[0] should be the first text assistant message, got %T %+v", msgs[0], msgs[0]) + } + if _, ok := msgs[1].(agentcore.UserMessage); !ok { + t.Errorf("msg[1] should be the injected follow-up user message, got %T", msgs[1]) + } + if a, ok := msgs[2].(agentcore.AssistantMessage); !ok || len(a.ToolCalls()) != 1 { + t.Errorf("msg[2] should be the tool-call assistant message, got %T %+v", msgs[2], msgs[2]) + } + tr, ok := msgs[3].(agentcore.ToolResultMessage) + if !ok || tr.ToolCallID != "call-1" || tr.IsError { + t.Errorf("msg[3] should be the successful echo tool result, got %T %+v", msgs[3], msgs[3]) + } + if a, ok := msgs[4].(agentcore.AssistantMessage); !ok || textContentOf(a.Content) != "all done" { + t.Errorf("msg[4] should be the final text assistant message, got %T %+v", msgs[4], msgs[4]) + } + + // The loop must have driven the provider exactly three times, each carrying + // the growing context and the configured model. + if p.callCount() != 3 { + t.Fatalf("provider called %d times, want 3", p.callCount()) + } + if req := p.requestAt(0); req.Model != "faux" { + t.Errorf("provider request model = %q, want faux", req.Model) + } +} + +// textContentOf returns the concatenated text of a content list. +func textContentOf(list agentcore.ContentList) string { + var s string + for _, c := range list { + if tc, ok := c.(agentcore.TextContent); ok { + s += tc.Text + } + } + return s +} + +// TestFauxSeamSixHooks exercises all six loop hooks through the real seam in a +// single run: the two per-request LoopConfig hooks (TransformContext, +// ConvertToLlm resolved via GetAPIKey) and the four RunConfig hooks +// (GetFollowUpMessages, GetSteeringMessages, PrepareNextTurn, +// ShouldStopAfterTurn). Each hook records that it fired and, where observable, +// that its effect reached the provider request. +func TestFauxSeamSixHooks(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{ + toolCallTurn("call-1", "echo", `{}`), // turn 1: tool → afterTurn hooks fire + textTurn("second"), // turn 2: end (after model swap) + }, + } + var fired struct { + transform, convert, apiKey, followUp, steering, prepare, shouldStop bool + } + swapped := "swapped-model" + cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false)) + cfg.Provider = "faux" + cfg.TransformContext = func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList { + fired.transform = true + return msgs + } + cfg.ConvertToLlm = func(msgs agentcore.MessageList) agentcore.MessageList { + fired.convert = true + return msgs + } + cfg.GetAPIKey = func(ctx context.Context, provider string) string { + fired.apiKey = true + return "dyn-key" + } + cfg.GetSteeringMessages = func(ctx context.Context) []agentcore.AgentMessage { + fired.steering = true + return nil + } + cfg.PrepareNextTurn = func(ctx context.Context, agentCtx *agentcore.AgentContext) *TurnUpdate { + fired.prepare = true + return &TurnUpdate{Model: &swapped} + } + stopCalls := 0 + cfg.ShouldStopAfterTurn = func(ctx context.Context, agentCtx *agentcore.AgentContext) bool { + fired.shouldStop = true + stopCalls++ + return false // never stop early; let the run end naturally + } + cfg.GetFollowUpMessages = func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage { + fired.followUp = true + // No follow-up: the tool-call turn already drives turn 2, so the run + // ends naturally after the second turn. + return nil + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}} + + collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + if !fired.transform || !fired.convert || !fired.apiKey { + t.Errorf("per-request hooks not all fired: %+v", fired) + } + if !fired.followUp || !fired.steering || !fired.prepare || !fired.shouldStop { + t.Errorf("per-turn hooks not all fired: %+v", fired) + } + if stopCalls == 0 { + t.Error("ShouldStopAfterTurn was never consulted") + } + // GetAPIKey's dynamic key must have reached the provider request config. + if got := p.requestAt(0).Config.APIKey; got != "dyn-key" { + t.Errorf("GetAPIKey result not threaded to provider, APIKey = %q", got) + } + // PrepareNextTurn swapped the model before turn 2. + if p.callCount() >= 2 { + if got := p.requestAt(1).Config.APIKey; got != "dyn-key" { + t.Errorf("turn 2 APIKey = %q, want dyn-key", got) + } + if got := p.requestAt(1).Model; got != swapped { + t.Errorf("PrepareNextTurn model swap not applied, turn 2 model = %q, want %q", got, swapped) + } + } +} + +// TestFauxSeamTruncationProtection verifies that a truncated (stopReason=length) +// tool-call turn is protected: the tool is NOT executed and a synthesized failed +// tool result is fed back, all through the seam. +func TestFauxSeamTruncationProtection(t *testing.T) { + // Turn 1: a tool call that arrives truncated. Turn 2: end. + truncPartial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewToolCallContent("t1", "echo", json.RawMessage(`{}`))}} + truncFinal := truncPartial + truncFinal.StopReason = agentcore.StopReasonLength + p := &fauxProvider{ + name: "faux", + turns: []fauxTurn{ + { + provider.StreamStartEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}}, + provider.StreamToolCallEvent{Partial: truncPartial}, + provider.StreamDoneEvent{Message: truncFinal}, + }, + textTurn("recovered"), + }, + } + cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false)) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + if countKind(kinds, agentcore.EventToolExecutionEnd) != 0 { + t.Errorf("truncated tool call must not execute, got %v", kinds) + } + var foundFail bool + for _, m := range msgs { + if tr, ok := m.(agentcore.ToolResultMessage); ok && tr.IsError && tr.ToolCallID == "t1" { + foundFail = true + } + } + if !foundFail { + t.Errorf("expected a synthesized failed tool result for the truncated call, got %+v", msgs) + } +} + +// TestFauxSeamParallelOrderingPreserved verifies that a turn with multiple +// parallel tool calls yields tool results in source order regardless of which +// tool finishes first, driven through the seam. +func TestFauxSeamParallelOrderingPreserved(t *testing.T) { + // One assistant turn with three tool calls in a fixed order; the tools sleep + // in reverse so completion order differs from source order. + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{ + agentcore.NewToolCallContent("a0", "slow", json.RawMessage(`{}`)), + agentcore.NewToolCallContent("a1", "mid", json.RawMessage(`{}`)), + agentcore.NewToolCallContent("a2", "fast", json.RawMessage(`{}`)), + }} + final := partial + final.StopReason = agentcore.StopReasonToolUse + p := &fauxProvider{ + turns: []fauxTurn{ + {provider.StreamStartEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}}, provider.StreamToolCallEvent{Partial: partial}, provider.StreamDoneEvent{Message: final}}, + textTurn("done"), + }, + } + mk := func(name string, delay time.Duration) execTool { + return execTool{ + name: name, + mode: agentcore.ToolExecutionParallel, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + time.Sleep(delay) + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } + } + cfg := newFauxRunCfg(p, mk("slow", 25*time.Millisecond), mk("mid", 12*time.Millisecond), mk("fast", 1*time.Millisecond)) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + _, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + var order []string + for _, m := range msgs { + if tr, ok := m.(agentcore.ToolResultMessage); ok { + order = append(order, tr.ToolCallID) + } + } + want := []string{"a0", "a1", "a2"} + if len(order) != 3 || order[0] != want[0] || order[1] != want[1] || order[2] != want[2] { + t.Errorf("parallel tool results out of source order: got %v, want %v", order, want) + } +} + +// TestFauxSeamStreamCancellation verifies that cancelling the context stops the +// run: the consumer stops receiving events and Result reports the cancellation, +// exercised through the seam with a provider that streams slowly. +func TestFauxSeamStreamCancellation(t *testing.T) { + p := &fauxProvider{ + turns: []fauxTurn{textTurn("never fully delivered")}, + delay: 50 * time.Millisecond, // slow enough to cancel mid-stream + } + cfg := newFauxRunCfg(p) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + ctx, cancel := context.WithCancel(context.Background()) + s := agentLoop(ctx, agentCtx, cfg) + + // Read the first event, then cancel while the provider is still streaming. + <-s.Events() + cancel() + // Drain remaining events (must terminate, not hang). + for range s.Events() { + } + if _, err := s.Result(context.Background()); err != nil { + // A set result is also acceptable (the run may have finished emitting + // agent_end before cancellation propagated); but if an error is set it + // must be the cancellation. + if err != context.Canceled { + t.Errorf("cancelled run result error = %v, want context.Canceled or nil", err) + } + } +} diff --git a/pigo/internal/runtime/headless.go b/pigo/internal/runtime/headless.go new file mode 100644 index 0000000..8c16f97 --- /dev/null +++ b/pigo/internal/runtime/headless.go @@ -0,0 +1,259 @@ +// This file implements the headless / stdio run modes (US-020, FR-18): a +// non-interactive driver that runs the agent loop over a single prompt for +// scripting and CI. Two output modes are supported, mirroring pi's print-mode +// and rpc/stream-json protocols: +// +// - PrintMode: run the loop to completion and write only the final assistant +// text to the output (the "-p / --print" mode). +// - StreamJSONMode: serialize every AgentEvent as a line-delimited JSON object +// as it is emitted (the "--output-format stream-json" mode), so a parent +// process can consume the run incrementally. +// +// The run's success/failure is reported as a returned error so the CLI can map +// it to a process exit code: a run whose final assistant message carries +// stopReason error/aborted, or whose stream result errors, is a failure. +package runtime + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// HeadlessMode selects how a headless run reports its progress and result. +type HeadlessMode int + +const ( + // PrintMode runs the loop to completion and writes only the final assistant + // text to the output writer. + PrintMode HeadlessMode = iota + // StreamJSONMode writes each AgentEvent as a line-delimited JSON object as it + // is emitted. + StreamJSONMode +) + +// HeadlessConfig configures a headless run. +type HeadlessConfig struct { + // Run is the loop configuration (provider stream, tools, hooks). + Run RunConfig + // Mode selects print vs stream-json output. Defaults to PrintMode. + Mode HeadlessMode + // Out receives the run output (final text or JSON lines). Required. + Out io.Writer + // OnEvent, when non-nil, is invoked for every AgentEvent before output + // handling. It is the seam plugin lifecycle-event delivery (US-017, #133) + // hooks into, independent of the output mode. It must not block. + OnEvent func(ev agentcore.AgentEvent) + // Progress receives human-readable sub-agent progress lines + // (SubAgentProgressEvent). Defaults to os.Stderr when nil. Progress is + // stderr-only by contract: it is never serialised onto Out (stdout), so it + // cannot pollute the final result text or the stream-json envelope stream. + Progress io.Writer +} + +// ErrRunFailed is the sentinel returned by RunHeadless when the agent run ended +// in a failure state (stopReason error/aborted). The CLI maps a non-nil error +// to a non-zero exit code. +type ErrRunFailed struct { + // Reason is the stopReason (or message) that marked the run as failed. + Reason string +} + +func (e *ErrRunFailed) Error() string { + if e.Reason == "" { + return "agent run failed" + } + return "agent run failed: " + e.Reason +} + +// RunHeadless runs the agent loop for the already-assembled agentCtx and drives +// output per cfg.Mode. It blocks until the run ends and returns nil on success +// or an error describing the failure (for exit-code mapping). It never returns +// before the stream is fully drained, so no goroutine is leaked. +func RunHeadless(ctx context.Context, agentCtx *agentcore.AgentContext, cfg HeadlessConfig) error { + if cfg.Out == nil { + return fmt.Errorf("headless: nil output writer") + } + stream := agentLoop(ctx, agentCtx, cfg.Run) + + // writeErr holds the first stream-json write failure. We keep draining after + // it (DrainStream never returns early) so the loop's producer goroutine never + // blocks on a synchronous Emit — honoring the no-leak contract on a broken + // pipe. In stream-json mode every event is serialised; print mode only needs + // the final message, which DrainStream returns. + var writeErr error + h := StreamHandler{} + // Compose the output-mode serialiser (stream-json only) with the optional + // external OnEvent (plugin lifecycle delivery, US-017). Both observe every + // event; the serialiser runs first so a write failure is recorded even when a + // plugin observer is also wired. + // + // SubAgentProgressEvent (D-9) is special-cased: it is written as a + // human-readable line to the progress writer (stderr) and is deliberately + // excluded from the stream-json stdout path so it never pollutes the result + // output or the machine-readable envelope stream (progress is stderr-only). + progress := cfg.Progress + if progress == nil { + progress = os.Stderr + } + streamJSON := cfg.Mode == StreamJSONMode + h.OnEvent = func(ev agentcore.AgentEvent) { + if pe, ok := ev.(agentcore.SubAgentProgressEvent); ok { + writeProgressLine(progress, pe) + if cfg.OnEvent != nil { + cfg.OnEvent(ev) + } + return + } + if streamJSON && writeErr == nil { + if err := writeEventJSON(cfg.Out, ev); err != nil { + writeErr = err + } + } + if cfg.OnEvent != nil { + cfg.OnEvent(ev) + } + } + lastAssistant, resErr := DrainStream(ctx, stream, h) + if writeErr != nil { + return writeErr + } + if resErr != nil { + return resErr + } + + if cfg.Mode == PrintMode { + text := "" + if lastAssistant != nil { + text = agentcore.ContentToText(lastAssistant.Content) + } + if _, err := io.WriteString(cfg.Out, text); err != nil { + return err + } + if text != "" && !strings.HasSuffix(text, "\n") { + if _, err := io.WriteString(cfg.Out, "\n"); err != nil { + return err + } + } + } + + if lastAssistant != nil { + switch lastAssistant.StopReason { + case agentcore.StopReasonError: + reason := lastAssistant.ErrorMessage + if reason == "" { + reason = "error" + } + return &ErrRunFailed{Reason: reason} + case agentcore.StopReasonAborted: + return &ErrRunFailed{Reason: "aborted"} + } + } + return nil +} + +// writeProgressLine renders one SubAgentProgressEvent as a human-readable line +// on w (stderr by contract). When the task supplied a description it is shown +// alongside the activity; otherwise the line degrades to the activity alone +// (Description MAY be empty, Activity never is). Write errors are ignored: +// progress is a non-critical, best-effort side channel. +func writeProgressLine(w io.Writer, ev agentcore.SubAgentProgressEvent) { + if ev.Description != "" { + fmt.Fprintf(w, " ⏺ %s · %s\n", ev.Description, ev.Activity) + return + } + fmt.Fprintf(w, " ⏺ %s\n", ev.Activity) +} + +// writeEventJSON serializes one AgentEvent as a single line of JSON, terminated +// by a newline, onto w. The envelope always carries a "type" discriminant so a +// consumer can dispatch without positional knowledge. +func writeEventJSON(w io.Writer, ev agentcore.AgentEvent) error { + env := eventEnvelope(ev) + b, err := json.Marshal(env) + if err != nil { + return fmt.Errorf("headless: marshal event: %w", err) + } + b = append(b, '\n') + _, err = w.Write(b) + return err +} + +// eventEnvelope maps an AgentEvent onto a JSON-serializable object with a +// "type" discriminant plus the event's observable payload. Only fields that are +// safe and useful over the wire are included (assistant text, tool ids/names, +// stop reasons) — never secrets. +func eventEnvelope(ev agentcore.AgentEvent) map[string]any { + env := map[string]any{"type": ev.EventType()} + switch e := ev.(type) { + case agentcore.AgentStartEvent: + // The first event carries the backing session id (mirrors pi/Claude Code), + // so a consumer can associate the run's output with a session and resume + // it later. Omitted only when the run has no backing session (SessionID + // unset), which the envelope treats as "not resumable". + if e.SessionID != "" { + env["sessionId"] = e.SessionID + } + case agentcore.AgentEndEvent: + env["messageCount"] = len(e.Messages) + case agentcore.TurnEndEvent: + env["stopReason"] = e.Message.StopReason + if text := agentcore.ContentToText(e.Message.Content); text != "" { + env["text"] = text + } + if calls := e.Message.ToolCalls(); len(calls) > 0 { + names := make([]string, len(calls)) + for i, c := range calls { + names[i] = c.Name + } + env["toolCalls"] = names + } + case agentcore.MessageUpdateEvent: + if a, ok := e.Message.(agentcore.AssistantMessage); ok { + if text := agentcore.ContentToText(a.Content); text != "" { + env["text"] = text + } + } + case agentcore.ToolExecutionStartEvent: + env["toolCallId"] = e.ToolCallID + env["toolName"] = e.ToolName + case agentcore.ToolExecutionEndEvent: + env["toolCallId"] = e.ToolCallID + env["toolName"] = e.ToolName + env["isError"] = e.IsError + case agentcore.CompactionStartEvent: + env["reason"] = e.Reason + env["tokensBefore"] = e.TokensBefore + case agentcore.CompactionEvent: + env["reason"] = e.Reason + env["tokensBefore"] = e.TokensBefore + env["tokensAfter"] = e.TokensAfter + env["summarizedCount"] = e.SummarizedCount + env["keptCount"] = e.KeptCount + if e.ErrorMessage != "" { + env["error"] = e.ErrorMessage + } + case agentcore.TelemetryEvent: + // The run-end telemetry summary: structured metrics a script can read + // directly from the stream-json output (observability -- structured telemetry collection). Per-tool + // timings are flattened into a name→{count,totalMs} object so a JSON + // consumer can index by tool name. + env["turns"] = e.Turns + env["truncationCount"] = e.TruncationCount + env["compactionCount"] = e.CompactionCount + env["contextUtilization"] = e.ContextUtilization + env["contextTokens"] = e.ContextTokens + env["contextWindow"] = e.ContextWindow + tools := make(map[string]map[string]any, len(e.ToolDurationsMs)) + for name, t := range e.ToolDurationsMs { + tools[name] = map[string]any{"count": t.Count, "totalMs": t.TotalMs} + } + env["toolDurationsMs"] = tools + } + return env +} diff --git a/pigo/internal/runtime/headless_test.go b/pigo/internal/runtime/headless_test.go new file mode 100644 index 0000000..136025e --- /dev/null +++ b/pigo/internal/runtime/headless_test.go @@ -0,0 +1,292 @@ +package runtime + +// This file is the end-to-end test for the headless / stdio run modes (US-020, +// #39). It drives RunHeadless over the real faux provider seam (no loop-internal +// mocking) and asserts the two output contracts — PrintMode's final text and +// StreamJSONMode's line-delimited JSON events — plus the success/failure signal +// that the CLI maps to a process exit code. + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// TestRunHeadlessPrintMode runs a text→tool→text scenario through RunHeadless in +// PrintMode and asserts that only the final assistant text reaches the writer, +// terminated by a newline, and that the run reports success (nil error). +func TestRunHeadlessPrintMode(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{ + toolCallTurn("call-1", "echo", `{"msg":"hi"}`), // turn 1: tool call + textTurn("final answer"), // turn 2: final text + }, + } + cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false)) + var out bytes.Buffer + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}} + + err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Mode: PrintMode, Out: &out}) + if err != nil { + t.Fatalf("RunHeadless print mode: unexpected error %v", err) + } + got := out.String() + if got != "final answer\n" { + t.Errorf("print mode output = %q, want %q", got, "final answer\n") + } +} + +// TestRunHeadlessStreamJSON runs the same scenario in StreamJSONMode and asserts +// every line is a valid JSON object carrying a "type" discriminant, that the run +// is bracketed by agent_start/agent_end, and that a tool execution is reported — +// the machine-readable protocol a parent process consumes. +func TestRunHeadlessStreamJSON(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{ + toolCallTurn("call-1", "echo", `{"msg":"hi"}`), + textTurn("done"), + }, + } + cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false)) + var out bytes.Buffer + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}} + + if err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Mode: StreamJSONMode, Out: &out}); err != nil { + t.Fatalf("RunHeadless stream-json: unexpected error %v", err) + } + + var types []string + sc := bufio.NewScanner(&out) + for sc.Scan() { + line := sc.Bytes() + if len(bytes.TrimSpace(line)) == 0 { + continue + } + var env map[string]any + if err := json.Unmarshal(line, &env); err != nil { + t.Fatalf("stream-json line is not valid JSON: %q (%v)", line, err) + } + typ, ok := env["type"].(string) + if !ok || typ == "" { + t.Errorf("stream-json line missing type discriminant: %q", line) + } + types = append(types, typ) + } + if len(types) == 0 { + t.Fatal("stream-json produced no event lines") + } + if types[0] != agentcore.EventAgentStart || types[len(types)-1] != agentcore.EventAgentEnd { + t.Errorf("stream must be bracketed by agent_start/agent_end, got %v", types) + } + if !contains(types, agentcore.EventToolExecutionEnd) { + t.Errorf("expected a tool_execution_end event, got %v", types) + } +} + +// TestRunHeadlessStreamJSONSessionID verifies that when RunConfig.SessionID is +// set, the first stream-json event (agent_start) carries it under "sessionId", +// so a consumer can associate the run's output with a session and resume it +// later (mirrors pi/Claude Code). When SessionID is empty the key is omitted. +func TestRunHeadlessStreamJSONSessionID(t *testing.T) { + run := func(sessionID string) map[string]any { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{textTurn("done")}, + } + cfg := newFauxRunCfg(p) + cfg.SessionID = sessionID + var out bytes.Buffer + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}} + if err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Mode: StreamJSONMode, Out: &out}); err != nil { + t.Fatalf("RunHeadless stream-json: unexpected error %v", err) + } + sc := bufio.NewScanner(&out) + for sc.Scan() { + line := sc.Bytes() + if len(bytes.TrimSpace(line)) == 0 { + continue + } + var env map[string]any + if err := json.Unmarshal(line, &env); err != nil { + t.Fatalf("stream-json line is not valid JSON: %q (%v)", line, err) + } + if env["type"] == agentcore.EventAgentStart { + return env + } + } + t.Fatal("no agent_start event found") + return nil + } + + first := run("sess-123") + if got, ok := first["sessionId"].(string); !ok || got != "sess-123" { + t.Errorf("agent_start sessionId = %v, want %q", first["sessionId"], "sess-123") + } + + none := run("") + if _, present := none["sessionId"]; present { + t.Errorf("agent_start must omit sessionId when SessionID is empty, got %v", none["sessionId"]) + } +} + +// TestRunHeadlessReportsFailure verifies that a run whose final assistant message +// carries stopReason=error surfaces as an ErrRunFailed, so the CLI maps it to a +// non-zero exit code. +func TestRunHeadlessReportsFailure(t *testing.T) { + errPartial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + errFinal := errPartial + errFinal.StopReason = agentcore.StopReasonError + errFinal.ErrorMessage = "boom" + p := &fauxProvider{ + name: "faux", + turns: []fauxTurn{ + { + provider.StreamStartEvent{Partial: errPartial}, + provider.StreamDoneEvent{Message: errFinal}, + }, + }, + } + cfg := newFauxRunCfg(p) + var out bytes.Buffer + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Mode: PrintMode, Out: &out}) + if err == nil { + t.Fatal("run ending in stopReason=error must return a non-nil error") + } + var failed *ErrRunFailed + if !as(err, &failed) { + t.Fatalf("error = %T (%v), want *ErrRunFailed", err, err) + } + if !strings.Contains(failed.Error(), "boom") { + t.Errorf("error message = %q, want it to mention the failure reason", failed.Error()) + } +} + +// TestRunHeadlessNilWriter guards the misconfiguration path. +func TestRunHeadlessNilWriter(t *testing.T) { + p := &fauxProvider{turns: []fauxTurn{textTurn("x")}} + cfg := newFauxRunCfg(p) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + if err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Out: nil}); err == nil { + t.Fatal("nil output writer must be rejected") + } +} + +// emitTool returns a tool that surfaces ev on the run stream via the run-level +// progress emitter the loop injects into ctx (WithProgressEmitter), then returns +// a trivial text result. This mirrors how a dispatched sub-agent surfaces a +// SubAgentProgressEvent up the parent stream. +func emitTool(name string, ev agentcore.AgentEvent) execTool { + return execTool{ + name: name, + mode: agentcore.ToolExecutionParallel, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + if emit := agentcore.ProgressEmitterFromContext(ctx); emit != nil { + _ = emit(ctx, ev) + } + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}}, nil + }, + } +} + +// TestRunHeadlessSubAgentProgressToStderr verifies the D-9 contract: a +// SubAgentProgressEvent emitted during the run is rendered as a human-readable +// line to the progress writer (stderr) and is NEVER serialised onto stdout — +// neither the final result text nor the stream-json envelope stream may contain +// it. The event is injected via a faux tool whose execution fires it on the +// run's event stream (the same seam the loop uses). +func TestRunHeadlessSubAgentProgressToStderr(t *testing.T) { + const desc = "investigate the parser" + const activity = "Editing" + + run := func(mode HeadlessMode) (stdout, stderr string) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{ + toolCallTurn("call-1", "task", `{"description":"investigate the parser"}`), + textTurn("done"), + }, + } + // The tool emits a SubAgentProgressEvent onto the run stream, mimicking a + // dispatched sub-agent surfacing progress up the parent stream. + tool := emitTool("task", agentcore.SubAgentProgressEvent{ + ToolCallID: "call-1", + Description: desc, + Activity: activity, + }) + cfg := newFauxRunCfg(p, tool) + var out, prog bytes.Buffer + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}} + if err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Mode: mode, Out: &out, Progress: &prog}); err != nil { + t.Fatalf("RunHeadless: unexpected error %v", err) + } + return out.String(), prog.String() + } + + for _, mode := range []struct { + name string + mode HeadlessMode + }{{"print", PrintMode}, {"stream-json", StreamJSONMode}} { + t.Run(mode.name, func(t *testing.T) { + stdout, stderr := run(mode.mode) + // (a) stderr carries the progress line with description + activity. + if !strings.Contains(stderr, desc) || !strings.Contains(stderr, activity) { + t.Errorf("stderr = %q, want it to contain description %q and activity %q", stderr, desc, activity) + } + // (b) stdout must not contain the progress event in any form. + if strings.Contains(stdout, "subagent_progress") { + t.Errorf("stdout must not contain the subagent_progress envelope, got %q", stdout) + } + if strings.Contains(stdout, desc) { + t.Errorf("stdout must not leak the progress description, got %q", stdout) + } + }) + } +} + +// TestWriteProgressLineEmptyDescription verifies the line degrades gracefully to +// the activity alone when the task supplied no description. +func TestWriteProgressLineEmptyDescription(t *testing.T) { + var buf bytes.Buffer + writeProgressLine(&buf, agentcore.SubAgentProgressEvent{Activity: "Thinking"}) + got := buf.String() + if !strings.Contains(got, "Thinking") { + t.Errorf("line = %q, want it to contain the activity", got) + } + if strings.Contains(got, "·") { + t.Errorf("line = %q, want no separator when description is empty", got) + } +} + +// contains reports whether s contains v. +func contains(s []string, v string) bool { + for _, x := range s { + if x == v { + return true + } + } + return false +} + +// as is a tiny errors.As shim kept local to avoid an extra import in a test that +// only ever unwraps one level. +func as(err error, target **ErrRunFailed) bool { + if e, ok := err.(*ErrRunFailed); ok { + *target = e + return true + } + return false +} diff --git a/pigo/internal/runtime/loop.go b/pigo/internal/runtime/loop.go new file mode 100644 index 0000000..f2be811 --- /dev/null +++ b/pigo/internal/runtime/loop.go @@ -0,0 +1,503 @@ +// This file implements pi's two-layer agent loop (US-006, FR-1). It strings +// together streaming assistant responses, batch tool execution, and the loop's +// six hooks with control flow kept faithful to pi's runLoop: +// +// - Inner loop: one turn = stream an assistant response → execute its tool +// calls → feed the results back, repeating until an assistant message has no +// tool calls (a natural turn end). +// - Outer loop: after the inner loop settles, pull getFollowUpMessages; if any +// are returned they become the next pending input and the inner loop runs +// again, otherwise the run ends. +// +// Per-turn hooks after each turn_end: getSteeringMessages (pulled after tool +// execution and injected before the next turn), prepareNextTurn (may swap +// context / model / thinkingLevel), shouldStopAfterTurn (true ⇒ agent_end + +// exit). Two stop reasons are handled specially: length (the response was +// truncated by the token cap) fails every tool call so the model resends +// (failToolCallsFromTruncatedMessage); error / aborted end the run immediately. +// +// agentLoop starts a fresh run from a prompt already appended to the context. +package runtime + +import ( + "context" + "fmt" + "os" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/provider" +) + +// nowMillis returns the current Unix time in milliseconds, the timestamp unit +// used for CompactionMessage checkpoints. +func nowMillis() int64 { return time.Now().UnixMilli() } + +// TurnUpdate is the optional result of PrepareNextTurn: any non-nil field +// replaces the corresponding piece of loop state before the next turn. It lets +// a caller swap the trimmed context, system prompt, tool set, model, or +// thinking level between turns (FR-6). +type TurnUpdate struct { + Messages *agentcore.MessageList + SystemPrompt *string + Tools *[]agentcore.AgentTool + Model *string + ThinkingLevel *agentcore.ThinkingLevel +} + +// StopDecision is the result of the OnStop seam. Block=true prevents the run +// from ending; Guidance, when non-empty, is appended as a user-role message to +// steer the forced continuation (the Stop / SubagentStop hook's reason). The +// zero value (Block=false) lets the run end. +type StopDecision struct { + Block bool + Guidance string +} + +// RunConfig is the full configuration for a loop run: the per-turn streaming +// config (embedded LoopConfig), the batch tool-execution config, and the four +// loop-level hooks. Every hook is optional (nil = default behavior). +type RunConfig struct { + LoopConfig + // Batch holds the tool registry and the prepare/before/after hooks used to + // execute each assistant message's tool calls. + Batch agenttool.BatchConfig + + // GetFollowUpMessages is consulted after the inner loop settles (an assistant + // message with no tool calls). Returning messages continues the outer loop + // with them as the next input; returning none ends the run (FR-9). + GetFollowUpMessages func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage + // GetSteeringMessages is pulled after each turn's tool execution and injected + // before the next turn (pi per-turn semantics, FR-8). + GetSteeringMessages func(ctx context.Context) []agentcore.AgentMessage + // PrepareNextTurn runs after each turn_end and may swap context / model / + // thinkingLevel for the next turn (FR-6). + PrepareNextTurn func(ctx context.Context, agentCtx *agentcore.AgentContext) *TurnUpdate + // ShouldStopAfterTurn runs after each turn_end; true ends the run with an + // agent_end event (FR-7). + ShouldStopAfterTurn func(ctx context.Context, agentCtx *agentcore.AgentContext) bool + + // OnStop, when set, is consulted right before the run would end naturally (no + // tool calls and no follow-up messages). Returning a decision with Block=true + // keeps the loop running: any Guidance is appended as a user-role message to + // steer the continued run (Stop / SubagentStop hooks, US-008/009, FR-10). The + // seam carries no loop-protection itself — the caller's decorator owns the + // consecutive-block counter and the FR-12 force-stop limit, so an ill-behaved + // hook cannot loop forever. nil or a non-blocking decision lets the run end. + OnStop func(ctx context.Context, agentCtx *agentcore.AgentContext) *StopDecision + + // Reminders holds the per-turn system-reminder providers (US-002, FR-1/FR-2). + // When non-empty, ephemeral messages are injected into each + // turn's LLM request through the existing TransformContext seam, so they never + // enter the persisted history. nil / empty = no injection. + Reminders *ReminderRegistry + + // EventBuffer is the buffer size of the emitted EventStream. 0 gives fully + // synchronous back-pressure (matching pi's awaited emit). + EventBuffer int + + // SessionID, when set, is carried in the run's agent_start event so a + // stream-json consumer sees the backing session id in the first event and can + // resume the run later (mirrors pi/Claude Code). It is also the session key + // under which auto-compaction checkpoints are persisted (see MemoryRoot). + SessionID string + + // MemoryRoot, when non-empty (together with SessionID), enables checkpoint + // persistence for the "infinite context" feature (#480/#481): after a + // successful auto-compaction the collapsed prefix's summary is written as a + // checkpoint under /sessions//checkpoint.md so a later + // run can reload it instead of replaying the whole transcript. It is left "" + // when persistent memory is disabled (memory.enabled=false), which fully + // disables checkpoint writing. A checkpoint write failure is non-fatal. + MemoryRoot string +} + +// LoopEventStream is the stream returned by the loop entry points: it carries +// AgentEvents and yields the messages newly produced during the run. +type LoopEventStream = agentcore.EventStream[agentcore.AgentEvent, []agentcore.AgentMessage] + +// agentLoop starts a fresh run. The caller has already appended the initiating +// user message(s) to agentCtx.Messages. It returns immediately with an +// EventStream; a producer goroutine drives the loop and closes the stream when +// the run ends. +func agentLoop(ctx context.Context, agentCtx *agentcore.AgentContext, cfg RunConfig) *LoopEventStream { + stream := agentcore.NewEventStream[agentcore.AgentEvent, []agentcore.AgentMessage](cfg.EventBuffer) + go runLoop(ctx, agentCtx, cfg, stream) + return stream +} + +// StartRun is the exported entry point for a fresh run, used by out-of-package +// drivers (the interactive REPL, US-022). It is a thin wrapper over agentLoop so +// the loop internals stay unexported while callers outside the package can +// still launch a run and consume its event stream. +func StartRun(ctx context.Context, agentCtx *agentcore.AgentContext, cfg RunConfig) *LoopEventStream { + return agentLoop(ctx, agentCtx, cfg) +} + +// runLoop is the producer: it drives the two-layer loop, emitting events onto +// stream and setting the stream result to the messages produced during the run. +func runLoop(ctx context.Context, agentCtx *agentcore.AgentContext, cfg RunConfig, stream *LoopEventStream) { + // Wire per-turn system-reminder injection (US-002) onto the TransformContext + // seam. Reminders are appended to the request-shaped copy only, so they stay + // ephemeral: never written back to agentCtx.Messages, never persisted, never + // swept into a compaction summary. + if !cfg.Reminders.Empty() { + cfg.TransformContext = cfg.Reminders.wrapTransform(cfg.TransformContext) + } + startIdx := len(agentCtx.Messages) + // tel accumulates structured telemetry (turn count, per-tool durations, + // truncation count, compaction count, latest context-utilization ratio) from + // the events emitted below, surfaced as a TelemetryEvent at run end. + tel := newTelemetry() + // newMessages returns the messages appended since the run began. + newMessages := func() []agentcore.AgentMessage { + if len(agentCtx.Messages) <= startIdx { + return nil + } + out := make([]agentcore.AgentMessage, len(agentCtx.Messages)-startIdx) + copy(out, agentCtx.Messages[startIdx:]) + return out + } + emit := func(ev agentcore.AgentEvent) error { + tel.observe(ev) + return stream.Emit(ctx, ev) + } + // emitFrom wraps the raw stream.Emit callback handed to streamAssistantResponse + // and ExecuteToolCalls so telemetry observes those events (message_* and + // tool_execution_*) too, without changing their signatures. + emitFrom := func(c context.Context, ev agentcore.AgentEvent) error { + tel.observe(ev) + return stream.Emit(c, ev) + } + + // finish emits the telemetry summary then agent_end (unless suppressed by a + // prior emit error), records the run result, and closes the stream exactly + // once. Telemetry is emitted first so a consumer sees the run's structured + // metrics immediately before the terminal event. + finish := func() { + _ = emit(tel.summary()) + msgs := newMessages() + _ = emit(agentcore.AgentEndEvent{Messages: msgs}) + stream.SetResult(msgs) + stream.Close() + } + + if err := emit(agentcore.AgentStartEvent{SessionID: cfg.SessionID}); err != nil { + finish() + return + } + + for { // outer loop: pending / follow-up messages + for { // inner loop: turns until no tool calls + if err := emit(agentcore.TurnStartEvent{}); err != nil { + finish() + return + } + + assistant, err := streamAssistantResponse(ctx, agentCtx, cfg.LoopConfig, emitFrom) + if err != nil { + // emit was cancelled mid-stream; end the run. + finish() + return + } + + switch assistant.StopReason { + case agentcore.StopReasonLength: + // Truncated by the token cap: fail every tool call so the model + // resends, then continue feeding back. + toolResults := failToolCallsFromTruncatedMessage(agentCtx, assistant) + if err := emit(agentcore.TurnEndEvent{Message: assistant, ToolResults: toolResults}); err != nil { + finish() + return + } + if afterTurn(ctx, agentCtx, &cfg, true, emit, tel) { + finish() + return + } + continue + case agentcore.StopReasonError, agentcore.StopReasonAborted: + // Terminal failure: emit the turn end and stop. + _ = emit(agentcore.TurnEndEvent{Message: assistant}) + finish() + return + } + + calls := toAgentToolCalls(assistant.ToolCalls()) + if len(calls) == 0 { + // Natural turn end: no tools to run. + if err := emit(agentcore.TurnEndEvent{Message: assistant}); err != nil { + finish() + return + } + if afterTurn(ctx, agentCtx, &cfg, false, emit, tel) { + finish() + return + } + break // exit inner loop → consult follow-up messages + } + + // Inject the run-level emitter into the context so tools (notably the + // generic task tool) can retrieve it via ProgressEmitterFromContext and + // surface a dispatched sub-agent's progress up this parent event stream. + // emitFrom feeds the parent stream and is run-scoped, so a child's + // SubAgentProgressEvent lands on the right run's stream. + toolCtx := agentcore.WithProgressEmitter(ctx, emitFrom) + toolResults, allTerminate := agenttool.ExecuteToolCalls(toolCtx, cfg.Batch, calls, emitFrom) + for _, tr := range toolResults { + agentCtx.Messages = append(agentCtx.Messages, tr) + } + if err := emit(agentcore.TurnEndEvent{Message: assistant, ToolResults: toolResults}); err != nil { + finish() + return + } + if allTerminate { + // Every tool asked to terminate the run. + finish() + return + } + if afterTurn(ctx, agentCtx, &cfg, true, emit, tel) { + finish() + return + } + // Feed the tool results back into the next turn. + } + + // Inner loop settled: consult follow-up messages. + if cfg.GetFollowUpMessages != nil { + if follow := cfg.GetFollowUpMessages(ctx, agentCtx); len(follow) > 0 { + agentCtx.Messages = append(agentCtx.Messages, follow...) + continue // outer loop with the follow-ups as new input + } + } + // Stop hook: the run is about to end naturally. A hook may block the end + // and force a continuation, feeding its guidance back as the next input + // (US-008/009, FR-10). The consecutive-block counter and force-stop limit + // (FR-12) live in the decorator behind OnStop, so this seam stays simple. + if cfg.OnStop != nil { + if dec := cfg.OnStop(ctx, agentCtx); dec != nil && dec.Block { + if dec.Guidance != "" { + agentCtx.Messages = append(agentCtx.Messages, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(dec.Guidance)}, + }) + } + continue // outer loop: keep the run alive + } + } + break + } + + finish() +} + +// afterTurn runs the per-turn hooks after a turn_end. When hadToolExecution is +// true it first pulls getSteeringMessages and injects them before the next turn +// (pi per-turn semantics). It then applies prepareNextTurn, runs auto-compaction +// when the context has outgrown its window, and finally consults +// shouldStopAfterTurn, returning true when the run should end. +func afterTurn(ctx context.Context, agentCtx *agentcore.AgentContext, cfg *RunConfig, hadToolExecution bool, emit func(agentcore.AgentEvent) error, tel *telemetry) (stop bool) { + if hadToolExecution && cfg.GetSteeringMessages != nil { + if steer := cfg.GetSteeringMessages(ctx); len(steer) > 0 { + agentCtx.Messages = append(agentCtx.Messages, steer...) + } + } + if cfg.PrepareNextTurn != nil { + if upd := cfg.PrepareNextTurn(ctx, agentCtx); upd != nil { + applyTurnUpdate(agentCtx, cfg, upd) + } + } + maybeAutoCompact(ctx, agentCtx, cfg, emit, tel) + // Record the latest context-utilization ratio once the turn has settled (after + // any compaction), so the telemetry summary reports the current used/window + // figure. This runs even when auto-compaction is disabled so utilization is + // still observable whenever the context window is known. + if tel != nil && cfg.ContextWindow > 0 { + tokens := compaction.EstimateContextTokens(agentCtx.Messages).Tokens + tel.recordContext(tokens, cfg.ContextWindow) + } + if cfg.ShouldStopAfterTurn != nil { + return cfg.ShouldStopAfterTurn(ctx, agentCtx) + } + return false +} + +// maybeAutoCompact checks whether the context has outgrown its usable window and, +// if so, compacts it in place and emits a CompactionEvent. Compaction is a no-op +// when disabled, when the context window is unknown (<= 0), or when usage is +// under threshold. A compaction failure is non-fatal: the original context is +// preserved and a CompactionEvent carrying ErrorMessage is emitted so the failure +// is observable without aborting the run (US-004). +func maybeAutoCompact(ctx context.Context, agentCtx *agentcore.AgentContext, cfg *RunConfig, emit func(agentcore.AgentEvent) error, tel *telemetry) { + if !cfg.Compaction.Enabled || cfg.ContextWindow <= 0 { + return + } + before := compaction.EstimateContextTokens(agentCtx.Messages).Tokens + // Record pre-compaction utilization so the ratio reflects the peak that + // triggered (or nearly triggered) compaction even when the summary is read + // mid-run. afterTurn overwrites it with the post-settle figure. + if tel != nil { + tel.recordContext(before, cfg.ContextWindow) + } + if !compaction.ShouldCompact(before, cfg.ContextWindow, cfg.Compaction) { + return + } + // Signal the start so a front-end can show an in-progress indicator while the + // summarization request (an LLM call that blocks the loop) is in flight. + _ = emit(agentcore.CompactionStartEvent{Reason: "threshold", TokensBefore: before}) + res, err := runCompaction(ctx, agentCtx.Messages, cfg) + kept := len(agentCtx.Messages) + if err != nil { + _ = emit(agentcore.CompactionEvent{ + Reason: "threshold", + TokensBefore: before, + TokensAfter: before, + KeptCount: kept, + ErrorMessage: err.Error(), + }) + return + } + if res == nil { + // Nothing to summarize (cut point left no prefix); leave context as-is. + return + } + // Persist a checkpoint of the collapsed prefix before rewriting the context so + // a later run can reload it (infinite context, #480/#481). It reuses the + // summary compaction just produced — no extra LLM call — and is best-effort: + // a write failure is logged and the run continues on the compacted context. + writeCompactionCheckpoint(ctx, agentCtx.Messages, res, cfg) + now := nowMillis() + rebuilt := res.RebuildContext(agentCtx.Messages, now) + summarized := len(agentCtx.Messages) - (len(rebuilt) - 1) + agentCtx.Messages = rebuilt + after := compaction.EstimateContextTokens(rebuilt).Tokens + _ = emit(agentcore.CompactionEvent{ + Reason: "threshold", + TokensBefore: before, + TokensAfter: after, + SummarizedCount: summarized, + KeptCount: len(rebuilt) - 1, + }) +} + +// runCompaction invokes compaction.Compact with the loop's summarization config, +// falling back to the primary Stream/Model when the summary-specific fields are +// unset. Compact derives the cut point from settings.KeepRecentTokens. +func runCompaction(ctx context.Context, msgs agentcore.MessageList, cfg *RunConfig) (*compaction.CompactionResult, error) { + stream := cfg.SummaryStream + if stream == nil { + stream = cfg.Stream + } + model := cfg.SummaryModel + if model.ID == "" { + model = provider.Model{Provider: cfg.Provider, ID: cfg.Model, ContextWindow: cfg.ContextWindow} + } + // Resolve the API key the same way the primary turn does (dynamic key wins, + // static APIKey is the fallback) so the summarization stream authenticates + // against auth-requiring providers instead of failing with "missing API key". + key := cfg.APIKey + if cfg.GetAPIKey != nil { + if dyn := cfg.GetAPIKey(ctx, cfg.Provider); dyn != "" { + key = dyn + } + } + scfg := provider.StreamConfig{APIKey: key, ThinkingLevel: cfg.ThinkingLevel} + return compaction.Compact(ctx, stream, model, msgs, cfg.Compaction, -1, nil, "", scfg) +} + +// writeCompactionCheckpoint persists the just-produced compaction summary as a +// session checkpoint so a later run can reload the collapsed prefix instead of +// replaying it (#480/#481). It is a no-op unless checkpoint persistence is wired +// (MemoryRoot and SessionID both set) — which is how memory.enabled=false keeps +// the whole subsystem inert. It reuses res.Summary (no extra summarization call) +// via BuildCheckpoint, tagging the checkpoint with the compaction cut point as +// its watermark. All failures are non-fatal: they are logged to stderr and the +// run continues on the compacted context (WriteCheckpoint's log-and-continue +// contract). +func writeCompactionCheckpoint(ctx context.Context, msgs agentcore.MessageList, res *compaction.CompactionResult, cfg *RunConfig) { + if cfg.MemoryRoot == "" || cfg.SessionID == "" || res == nil { + return + } + watermark := res.FirstKeptIndex + if watermark < 0 { + watermark = 0 + } + if watermark > len(msgs) { + watermark = len(msgs) + } + // summarize returns the summary the compaction already computed, so + // BuildCheckpoint records an honest CoveredMessages count without a second + // LLM round-trip. + summarize := func(context.Context, []agentcore.Message) (string, error) { + return res.Summary, nil + } + cp, err := BuildCheckpoint(ctx, msgs[:watermark], watermark, time.Now(), summarize) + if err != nil { + fmt.Fprintf(os.Stderr, "pigo: checkpoint: build for session %s: %v\n", cfg.SessionID, err) + return + } + if err := WriteCheckpoint(cfg.SessionID, cfg.MemoryRoot, cp); err != nil { + fmt.Fprintf(os.Stderr, "pigo: checkpoint: write for session %s: %v\n", cfg.SessionID, err) + } +} + +// applyTurnUpdate applies a non-nil TurnUpdate to the mutable loop state: any +// set field replaces the current context / config value for the next turn. +func applyTurnUpdate(agentCtx *agentcore.AgentContext, cfg *RunConfig, upd *TurnUpdate) { + if upd.Messages != nil { + agentCtx.Messages = *upd.Messages + } + if upd.SystemPrompt != nil { + agentCtx.SystemPrompt = *upd.SystemPrompt + } + if upd.Tools != nil { + agentCtx.Tools = *upd.Tools + } + if upd.Model != nil { + cfg.Model = *upd.Model + } + if upd.ThinkingLevel != nil { + cfg.ThinkingLevel = *upd.ThinkingLevel + } +} + +// failToolCallsFromTruncatedMessage produces an error tool-result message for +// every tool call in a truncated (stopReason=length) assistant message, telling +// the model the response was cut off and to resend. The results are appended to +// the context and returned. Mirrors pi's failToolCallsFromTruncatedMessage. +func failToolCallsFromTruncatedMessage(agentCtx *agentcore.AgentContext, assistant agentcore.AssistantMessage) []agentcore.ToolResultMessage { + calls := assistant.ToolCalls() + if len(calls) == 0 { + return nil + } + results := make([]agentcore.ToolResultMessage, 0, len(calls)) + for _, c := range calls { + results = append(results, agentcore.ToolResultMessage{ + RoleField: agentcore.RoleToolResult, + ToolCallID: c.ID, + ToolName: c.Name, + Content: agentcore.ContentList{agentcore.NewTextContent( + "The previous response was truncated because it hit the output token limit, " + + "so this tool call was not executed. Please send a shorter response and retry.")}, + IsError: true, + }) + } + for _, r := range results { + agentCtx.Messages = append(agentCtx.Messages, r) + } + return results +} + +// toAgentToolCalls converts the assistant message's ToolCallContent blocks into +// the loop-level AgentToolCall view executeToolCalls consumes. +func toAgentToolCalls(blocks []agentcore.ToolCallContent) []agentcore.AgentToolCall { + if len(blocks) == 0 { + return nil + } + calls := make([]agentcore.AgentToolCall, len(blocks)) + for i, b := range blocks { + calls[i] = agentcore.AgentToolCall{ID: b.ID, Name: b.Name, Arguments: b.Arguments} + } + return calls +} diff --git a/pigo/internal/runtime/loop_onstop_test.go b/pigo/internal/runtime/loop_onstop_test.go new file mode 100644 index 0000000..1979fd3 --- /dev/null +++ b/pigo/internal/runtime/loop_onstop_test.go @@ -0,0 +1,63 @@ +package runtime + +import ( + "context" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// TestAgentLoopOnStopBlocksThenAllows: OnStop blocks the natural end twice +// (injecting guidance each time) then allows it, so the run takes three turns +// and the guidance messages land in the context. +func TestAgentLoopOnStopBlocksThenAllows(t *testing.T) { + cfg := newRunCfg(scriptedStream(nil)) // always a natural end_turn + blocks := 0 + cfg.OnStop = func(ctx context.Context, agentCtx *agentcore.AgentContext) *StopDecision { + if blocks < 2 { + blocks++ + return &StopDecision{Block: true, Guidance: "keep going"} + } + return nil + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, _ := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + if got := countKind(kinds, agentcore.EventTurnStart); got != 3 { + t.Fatalf("expected 3 turns (2 forced continuations + final), got %d (%v)", got, kinds) + } + if kinds[len(kinds)-1] != agentcore.EventAgentEnd { + t.Fatalf("run must end with agent_end, got %v", kinds) + } + guidance := 0 + for _, m := range agentCtx.Messages { + if um, ok := m.(agentcore.UserMessage); ok && len(um.Content) == 1 { + if tc, ok := um.Content[0].(agentcore.TextContent); ok && tc.Text == "keep going" { + guidance++ + } + } + } + if guidance != 2 { + t.Fatalf("expected 2 injected guidance messages, got %d", guidance) + } +} + +// TestAgentLoopOnStopNilEndsRun: a nil OnStop decision lets the run end after a +// single natural turn. +func TestAgentLoopOnStopNilEndsRun(t *testing.T) { + cfg := newRunCfg(scriptedStream(nil)) + consulted := false + cfg.OnStop = func(ctx context.Context, agentCtx *agentcore.AgentContext) *StopDecision { + consulted = true + return nil + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, _ := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + if !consulted { + t.Fatal("OnStop was never consulted") + } + if got := countKind(kinds, agentcore.EventTurnStart); got != 1 { + t.Fatalf("nil decision must end after one turn, got %d turns", got) + } +} diff --git a/pigo/internal/runtime/loop_test.go b/pigo/internal/runtime/loop_test.go new file mode 100644 index 0000000..f2127fa --- /dev/null +++ b/pigo/internal/runtime/loop_test.go @@ -0,0 +1,286 @@ +package runtime + +import ( + "context" + "encoding/json" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/provider" +) + +// collectStream drains a LoopEventStream, returning the event types in order +// and the run result messages. +func collectStream(t *testing.T, s *LoopEventStream) ([]string, []agentcore.AgentMessage) { + t.Helper() + var kinds []string + for ev := range s.Events() { + kinds = append(kinds, ev.EventType()) + } + msgs, err := s.Result(context.Background()) + if err != nil { + t.Fatalf("stream result: %v", err) + } + return kinds, msgs +} + +// oneToolAssistant builds an assistant message with a single tool call. +func oneToolAssistant(id, name string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + StopReason: agentcore.StopReasonToolUse, + Content: agentcore.ContentList{agentcore.NewToolCallContent(id, name, json.RawMessage(`{}`))}, + } +} + +// scriptedStream returns a StreamFn that emits one StreamDoneEvent per call, +// consuming msgs in order. Extra calls beyond msgs emit a plain end_turn. +func scriptedStream(msgs []agentcore.AssistantMessage) provider.StreamFn { + i := 0 + return func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + var msg agentcore.AssistantMessage + if i < len(msgs) { + msg = msgs[i] + } else { + msg = agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn} + } + i++ + s := provider.NewAssistantMessageEventStream(0) + go func() { + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: msg}) + s.Close() + }() + return s, nil + } +} + +func newRunCfg(stream provider.StreamFn, tools ...agentcore.AgentTool) RunConfig { + reg := agenttool.NewToolRegistry() + for _, tl := range tools { + _ = reg.Register(tl) + } + return RunConfig{ + LoopConfig: LoopConfig{Model: "fake", Stream: stream}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}}, + } +} + +func TestAgentLoopNoToolCallsSingleTurn(t *testing.T) { + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}, + })) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + want := []string{agentcore.EventAgentStart, agentcore.EventTurnStart, agentcore.EventMessageEnd, agentcore.EventTurnEnd, agentcore.EventTelemetry, agentcore.EventAgentEnd} + assertEventKinds(t, kinds, want) + if len(msgs) != 1 { + t.Fatalf("run produced %d messages, want 1: %+v", len(msgs), msgs) + } +} + +func TestAgentLoopInnerLoopFeedsToolResults(t *testing.T) { + // Turn 1: tool call. Turn 2: no tool call → inner loop ends. + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + oneToolAssistant("c1", "echo"), + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("done")}}, + }), echoTool("echo", agentcore.ToolExecutionParallel, false)) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + // Two turns; a tool executed in the first. + if countKind(kinds, agentcore.EventTurnStart) != 2 { + t.Errorf("expected 2 turns, got kinds %v", kinds) + } + if countKind(kinds, agentcore.EventToolExecutionEnd) != 1 { + t.Errorf("expected 1 tool execution, got kinds %v", kinds) + } + // Messages produced: assistant(tool) + toolResult + assistant(done) = 3. + if len(msgs) != 3 { + t.Fatalf("expected 3 new messages, got %d: %+v", len(msgs), msgs) + } + if _, ok := msgs[1].(agentcore.ToolResultMessage); !ok { + t.Errorf("expected message[1] to be a tool result, got %T", msgs[1]) + } +} + +func TestAgentLoopFollowUpMessagesContinue(t *testing.T) { + served := false + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("first")}}, + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("second")}}, + })) + cfg.GetFollowUpMessages = func(ctx context.Context, agentCtx *agentcore.AgentContext) []agentcore.AgentMessage { + if served { + return nil + } + served = true + return []agentcore.AgentMessage{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("more")}}} + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, _ := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + if countKind(kinds, agentcore.EventTurnStart) != 2 { + t.Errorf("follow-up should drive a second turn, got kinds %v", kinds) + } +} + +func TestAgentLoopShouldStopAfterTurn(t *testing.T) { + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + oneToolAssistant("c1", "echo"), + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}, + }), echoTool("echo", agentcore.ToolExecutionParallel, false)) + cfg.ShouldStopAfterTurn = func(ctx context.Context, agentCtx *agentcore.AgentContext) bool { return true } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, _ := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + // Stops after the first turn_end, so only one turn. + if countKind(kinds, agentcore.EventTurnStart) != 1 { + t.Errorf("shouldStopAfterTurn=true must stop after one turn, got %v", kinds) + } + if kinds[len(kinds)-1] != agentcore.EventAgentEnd { + t.Errorf("run must end with agent_end, got %v", kinds) + } +} + +func TestAgentLoopSteeringInjected(t *testing.T) { + var injectedSeen bool + steer := agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("steer")}} + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + oneToolAssistant("c1", "echo"), + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}, + }), echoTool("echo", agentcore.ToolExecutionParallel, false)) + pulled := false + cfg.GetSteeringMessages = func(ctx context.Context) []agentcore.AgentMessage { + if pulled { + return nil + } + pulled = true + return []agentcore.AgentMessage{steer} + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + for _, m := range agentCtx.Messages { + if um, ok := m.(agentcore.UserMessage); ok && len(um.Content) == 1 { + if tc, ok := um.Content[0].(agentcore.TextContent); ok && tc.Text == "steer" { + injectedSeen = true + } + } + } + if !injectedSeen { + t.Errorf("steering message was not injected into the context") + } +} + +func TestAgentLoopPrepareNextTurnSwapsModel(t *testing.T) { + var seenModels []string + streamFn := func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + seenModels = append(seenModels, model) + var msg agentcore.AssistantMessage + if len(seenModels) == 1 { + msg = oneToolAssistant("c1", "echo") + } else { + msg = agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn} + } + s := provider.NewAssistantMessageEventStream(0) + go func() { _ = s.Emit(ctx, provider.StreamDoneEvent{Message: msg}); s.Close() }() + return s, nil + } + cfg := newRunCfg(streamFn, echoTool("echo", agentcore.ToolExecutionParallel, false)) + newModel := "swapped-model" + cfg.PrepareNextTurn = func(ctx context.Context, agentCtx *agentcore.AgentContext) *TurnUpdate { + return &TurnUpdate{Model: &newModel} + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + if len(seenModels) != 2 || seenModels[1] != newModel { + t.Errorf("prepareNextTurn should swap model to %q, saw %v", newModel, seenModels) + } +} + +func TestAgentLoopLengthFailsToolCalls(t *testing.T) { + // Turn 1: tool call but truncated (length). Turn 2: end. + truncated := oneToolAssistant("c1", "echo") + truncated.StopReason = agentcore.StopReasonLength + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + truncated, + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}, + }), echoTool("echo", agentcore.ToolExecutionParallel, false)) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + // The tool must NOT have executed (truncated → failed instead). + if countKind(kinds, agentcore.EventToolExecutionEnd) != 0 { + t.Errorf("truncated message must not execute tools, got %v", kinds) + } + // A failed tool result must have been synthesized. + var foundFail bool + for _, m := range msgs { + if tr, ok := m.(agentcore.ToolResultMessage); ok && tr.IsError && tr.ToolCallID == "c1" { + foundFail = true + } + } + if !foundFail { + t.Errorf("expected a synthesized failed tool result for the truncated call") + } +} + +func TestAgentLoopErrorStopEndsRun(t *testing.T) { + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonError, ErrorMessage: "boom"}, + })) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, _ := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + if countKind(kinds, agentcore.EventTurnStart) != 1 { + t.Errorf("error stop must end after one turn, got %v", kinds) + } + if kinds[len(kinds)-1] != agentcore.EventAgentEnd { + t.Errorf("run must end with agent_end, got %v", kinds) + } +} + +func TestAgentLoopAllTerminateStopsRun(t *testing.T) { + term := true + termTool := execTool{ + name: "quit", + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("bye")}, Terminate: &term}, nil + }, + } + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + oneToolAssistant("c1", "quit"), + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}, // should never be reached + }), termTool) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + kinds, _ := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + if countKind(kinds, agentcore.EventTurnStart) != 1 { + t.Errorf("terminate must end the run after one turn, got %v", kinds) + } +} + +func assertEventKinds(t *testing.T, got, want []string) { + t.Helper() + if len(got) != len(want) { + t.Fatalf("event kinds = %v, want %v", got, want) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("event[%d] = %q, want %q (full %v)", i, got[i], want[i], got) + } + } +} + +func countKind(kinds []string, want string) int { + n := 0 + for _, k := range kinds { + if k == want { + n++ + } + } + return n +} diff --git a/pigo/internal/runtime/memory_reminder.go b/pigo/internal/runtime/memory_reminder.go new file mode 100644 index 0000000..bcfdf95 --- /dev/null +++ b/pigo/internal/runtime/memory_reminder.go @@ -0,0 +1,185 @@ +package runtime + +import ( + "context" + "sort" + "strings" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/memory" +) + +// defaultMemoryReminderMaxChars is the per-turn character budget for the +// injected memory body. It is deliberately modest so relevant memory context +// does not crowd out the live conversation in the model's window. +const defaultMemoryReminderMaxChars = 600 + +// defaultMemoryReminderLimit is the number of search hits fetched per turn +// before budget trimming. +const defaultMemoryReminderLimit = 5 + +// MemoryReminderProvider injects relevant persistent memory as ephemeral +// background context on each turn (issue #478). It derives a query from the +// most recent user message, runs a BM25 search over the memory store, and +// surfaces the top-ranked snippets — budget-capped and deduped so identical +// context is never repeated turn after turn. +// +// The returned body is RAW plain text; the reminder registry wraps it in +// tags and injects it only into the per-turn LLM request, so +// it never enters persisted history. +type MemoryReminderProvider struct { + // Store is the persistent memory database. When nil the provider never + // fires. + Store *memory.Store + + // MaxChars caps the injected body length. <=0 uses + // defaultMemoryReminderMaxChars. + MaxChars int + + // Limit caps the number of search hits considered. <=0 uses + // defaultMemoryReminderLimit. + Limit int + + // Scope and ScopeID, when non-empty, focus the search on a single memory + // scope (e.g. the current project id) via SearchOptions. + Scope string + ScopeID string + + // mu guards lastBody so Reminder is safe to call across turns. + mu sync.Mutex + lastBody string +} + +// Name implements ReminderProvider. +func (p *MemoryReminderProvider) Name() string { return "memory" } + +// Reminder implements ReminderProvider. It fires when the latest user message +// yields a search query that matches stored memory, returning a concise, +// budget-capped body of ranked snippets. It stays silent when the store is nil, +// there is no user text to search on, the search errors or returns nothing, or +// the produced body is identical to the one injected on the previous firing +// turn (dedupe). +func (p *MemoryReminderProvider) Reminder(ctx context.Context, msgs agentcore.MessageList) (string, bool) { + if p == nil || p.Store == nil { + return "", false + } + + query := latestUserText(msgs) + if strings.TrimSpace(query) == "" { + return "", false + } + + limit := p.Limit + if limit <= 0 { + limit = defaultMemoryReminderLimit + } + + results, err := p.Store.Search(query, memory.SearchOptions{ + Scope: p.Scope, + ScopeID: p.ScopeID, + Limit: limit, + ReconcileFirst: true, + }) + if err != nil || len(results) == 0 { + return "", false + } + + body := p.buildBody(results) + if body == "" { + return "", false + } + + // Dedupe: never re-inject the identical body on a subsequent firing turn. + p.mu.Lock() + defer p.mu.Unlock() + if body == p.lastBody { + return "", false + } + p.lastBody = body + return body, true +} + +// buildBody renders the ranked results into a concise, budget-capped body. +// MEMORY.md index files (and free-type entries) are ordered first as they are +// the most useful high-level context. +func (p *MemoryReminderProvider) buildBody(results []memory.SearchResult) string { + maxChars := p.MaxChars + if maxChars <= 0 { + maxChars = defaultMemoryReminderMaxChars + } + + // Stable sort so index-type / MEMORY.md hits float to the top while + // preserving the underlying BM25 order among equals. + ordered := make([]memory.SearchResult, len(results)) + copy(ordered, results) + sort.SliceStable(ordered, func(i, j int) bool { + return memoryRank(ordered[i]) < memoryRank(ordered[j]) + }) + + const heading = "Relevant memory:" + var b strings.Builder + b.WriteString(heading) + for _, r := range ordered { + snippet := strings.TrimSpace(strings.ReplaceAll(r.Snippet, "\n", " ")) + if snippet == "" { + continue + } + line := "\n- " + r.Path + ": " + snippet + // Enforce the budget: stop before exceeding maxChars, and never emit a + // heading-only body. + if b.Len()+len(line) > maxChars { + if b.Len() > len(heading) { + break + } + // The very first line already overflows: hard-truncate it so we + // still surface something within budget. + room := maxChars - b.Len() + if room <= 0 { + break + } + if room < len(line) { + line = line[:room] + } + b.WriteString(line) + break + } + b.WriteString(line) + } + + if b.Len() <= len(heading) { + return "" + } + return b.String() +} + +// memoryRank returns a sort key that floats MEMORY.md index files and free-type +// entries to the front (rank 0) ahead of everything else (rank 1). +func memoryRank(r memory.SearchResult) int { + if strings.HasSuffix(r.Path, "MEMORY.md") || r.Type == memory.TypeFree { + return 0 + } + return 1 +} + +// latestUserText returns the flattened text of the most recent user message in +// msgs, or "" if there is none. Reminder messages are user-role too, but they +// carry the preamble; those are skipped so the search query +// reflects the genuine user request rather than previously injected context. +func latestUserText(msgs agentcore.MessageList) string { + for i := len(msgs) - 1; i >= 0; i-- { + um, ok := msgs[i].(agentcore.UserMessage) + if !ok { + continue + } + text := agentcore.ContentToText(um.Content) + if strings.Contains(text, "") { + continue + } + if strings.TrimSpace(text) == "" { + continue + } + return text + } + return "" +} diff --git a/pigo/internal/runtime/memory_reminder_test.go b/pigo/internal/runtime/memory_reminder_test.go new file mode 100644 index 0000000..72d28cc --- /dev/null +++ b/pigo/internal/runtime/memory_reminder_test.go @@ -0,0 +1,150 @@ +package runtime + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/memory" +) + +// openMemoryStore opens a memory.Store over a temp DB + root and writes the +// given memory files (path segments relative to root -> body), returning the +// store. Reconcile is left to the provider's ReconcileFirst. +func openMemoryStore(t *testing.T, files map[string]string) *memory.Store { + t.Helper() + base := t.TempDir() + root := filepath.Join(base, "mimo") + if err := os.MkdirAll(root, 0o755); err != nil { + t.Fatalf("mkdir root: %v", err) + } + for rel, body := range files { + full := filepath.Join(root, rel) + if err := os.MkdirAll(filepath.Dir(full), 0o755); err != nil { + t.Fatalf("mkdir for %q: %v", full, err) + } + if err := os.WriteFile(full, []byte(body), 0o644); err != nil { + t.Fatalf("write %q: %v", full, err) + } + } + st, err := memory.Open(filepath.Join(base, "memory.db"), root, "") + if err != nil { + t.Fatalf("memory.Open: %v", err) + } + t.Cleanup(func() { st.Close() }) + return st +} + +// userMsgs builds a MessageList with a single user text message. +func userMsgs(text string) agentcore.MessageList { + return agentcore.MessageList{ + agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + }, + } +} + +func TestMemoryReminderInjectsMatchingSnippet(t *testing.T) { + st := openMemoryStore(t, map[string]string{ + filepath.Join("projects", "proj1", "notes", "auth.md"): "permission deadlock encountered during checkpoint save then retry succeeded", + filepath.Join("global", "user", "u1.md"): "unrelated grocery shopping list", + }) + p := &MemoryReminderProvider{Store: st, MaxChars: 400} + + body, ok := p.Reminder(context.Background(), userMsgs("how do I handle the permission deadlock?")) + if !ok { + t.Fatalf("expected a memory reminder to fire, got ok=false") + } + if !strings.Contains(body, "Relevant memory:") { + t.Errorf("body missing heading: %q", body) + } + if !strings.Contains(body, "permission") { + t.Errorf("body missing the matching snippet text: %q", body) + } + if len(body) > 400 { + t.Errorf("body exceeds MaxChars budget: len=%d body=%q", len(body), body) + } + + // Second identical call dedupes. + if _, ok := p.Reminder(context.Background(), userMsgs("how do I handle the permission deadlock?")); ok { + t.Errorf("identical follow-up call should dedupe to ok=false") + } +} + +func TestMemoryReminderRespectsMaxChars(t *testing.T) { + long := strings.Repeat("permission deadlock retry ", 40) + st := openMemoryStore(t, map[string]string{ + filepath.Join("projects", "proj1", "notes", "big.md"): long, + }) + p := &MemoryReminderProvider{Store: st, MaxChars: 120} + + body, ok := p.Reminder(context.Background(), userMsgs("permission deadlock")) + if !ok { + t.Fatalf("expected a memory reminder to fire") + } + if len(body) > 120 { + t.Errorf("body exceeds MaxChars=120: len=%d", len(body)) + } +} + +func TestMemoryReminderNoUserMessage(t *testing.T) { + st := openMemoryStore(t, map[string]string{ + filepath.Join("global", "user", "u1.md"): "permission deadlock note", + }) + p := &MemoryReminderProvider{Store: st} + + // Empty message list. + if _, ok := p.Reminder(context.Background(), nil); ok { + t.Errorf("empty message list must not fire") + } + // User message with no text content. + blank := agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}} + if _, ok := p.Reminder(context.Background(), blank); ok { + t.Errorf("empty user text must not fire") + } +} + +func TestMemoryReminderNoMatch(t *testing.T) { + st := openMemoryStore(t, map[string]string{ + filepath.Join("global", "user", "u1.md"): "grocery shopping list milk eggs", + }) + p := &MemoryReminderProvider{Store: st} + if _, ok := p.Reminder(context.Background(), userMsgs("kubernetes ingress controller crash")); ok { + t.Errorf("no matching memory should not fire") + } +} + +func TestMemoryReminderNilStore(t *testing.T) { + p := &MemoryReminderProvider{} + if _, ok := p.Reminder(context.Background(), userMsgs("anything")); ok { + t.Errorf("nil store must not fire") + } + var np *MemoryReminderProvider + if _, ok := np.Reminder(context.Background(), userMsgs("anything")); ok { + t.Errorf("nil provider must not fire") + } +} + +func TestMemoryReminderMemoryMdFirst(t *testing.T) { + st := openMemoryStore(t, map[string]string{ + filepath.Join("projects", "proj1", "notes", "detail.md"): "permission deadlock detail note here", + filepath.Join("projects", "proj1", "MEMORY.md"): "permission deadlock index overview", + }) + p := &MemoryReminderProvider{Store: st, MaxChars: 800} + body, ok := p.Reminder(context.Background(), userMsgs("permission deadlock")) + if !ok { + t.Fatalf("expected reminder to fire") + } + idxMem := strings.Index(body, "MEMORY.md") + idxDetail := strings.Index(body, "detail.md") + if idxMem == -1 || idxDetail == -1 { + t.Fatalf("expected both files in body: %q", body) + } + if idxMem > idxDetail { + t.Errorf("MEMORY.md should sort before other results, got body: %q", body) + } +} diff --git a/pigo/internal/runtime/orchestration_test.go b/pigo/internal/runtime/orchestration_test.go new file mode 100644 index 0000000..6c3371b --- /dev/null +++ b/pigo/internal/runtime/orchestration_test.go @@ -0,0 +1,638 @@ +package runtime + +// Tests for sub-agent orchestration, skills, and slash-commands (US-027/028/029, +// #45). The sub-agent integration test drives the flagship parent→child→parent +// path through the faux provider seam (mirrors faux_provider_test.go): the parent +// loop calls a sub-agent tool, the child runs its own scripted loop, and the +// child's final text is fed back as the parent's tool result. Skills and +// slash-commands get focused load/parse unit tests. + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/provider" +) + +// TestSubAgentParentChildParent is acceptance-critical: a parent agent loop +// delegates to a sub-agent via a tool call, the child runs an independent loop +// with its own context and provider, and the child's final assistant text is +// returned to the parent as the tool result (parent->child->parent). Both loops are driven by +// faux providers over the real StreamFnFromProvider seam. +func TestSubAgentParentChildParent(t *testing.T) { + // Child provider: a single scripted turn producing the delegated answer. + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{textTurn("child result: 42")}, + } + // The sub-agent tool spawns a child loop over its own context + provider. + sub := NewSubAgentTool(SubAgentSpec{ + Name: "researcher", + Description: "delegate research to a fresh sub-agent", + SystemPrompt: "you are a researcher", + Tools: nil, + NewRunConfig: func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: provider.StreamFnFromProvider(child)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: agenttool.NewToolRegistry()}}, + } + }, + }) + + // Parent provider: turn 1 calls the sub-agent, turn 2 (after the tool result + // is fed back) produces the final answer. + parent := &fauxProvider{ + name: "faux-parent", + models: []provider.Model{{Provider: "faux-parent", ID: "parent"}}, + turns: []fauxTurn{ + toolCallTurn("call-sub", "researcher", `{"prompt":"find the answer"}`), + textTurn("final: incorporated child result"), + }, + } + cfg := newFauxRunCfg(parent, sub) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("delegate this")}}, + }} + + kinds, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + // The sub-agent tool must have executed exactly once in the parent loop. + if got := countKind(kinds, agentcore.EventToolExecutionStart); got != 1 { + t.Errorf("expected 1 sub-agent tool execution, got %d in %v", got, kinds) + } + // The child provider must have been driven independently. + if child.callCount() != 1 { + t.Errorf("child provider called %d times, want 1", child.callCount()) + } + // The tool result fed back to the parent must carry the child's final text. + var toolResult *agentcore.ToolResultMessage + for i := range msgs { + if tr, ok := msgs[i].(agentcore.ToolResultMessage); ok { + toolResult = &tr + break + } + } + if toolResult == nil { + t.Fatalf("no tool result message in parent transcript: %+v", msgs) + } + if got := textContentOf(toolResult.Content); got != "child result: 42" { + t.Errorf("sub-agent result = %q, want %q (child final text fed to parent)", got, "child result: 42") + } + if toolResult.IsError { + t.Errorf("sub-agent tool result should not be an error") + } + // The parent's final message incorporates the delegation. + final := agentcore.LastAssistantOf(msgs) + if final == nil || textContentOf(final.Content) != "final: incorporated child result" { + t.Errorf("parent final text = %+v, want the post-delegation answer", final) + } +} + +// TestSubAgentConcurrent verifies multiple sub-agents can run concurrently: +// the parent issues two parallel sub-agent tool calls in one turn, each spawning +// an independent child loop, and both results are fed back. +func TestSubAgentConcurrent(t *testing.T) { + mkChild := func(answer string) *SubAgentTool { + cp := &fauxProvider{ + name: "faux-child-" + answer, + models: []provider.Model{{Provider: "c", ID: "c"}}, + turns: []fauxTurn{textTurn(answer)}, + } + return NewSubAgentTool(SubAgentSpec{ + Name: "agent-" + answer, + Description: "child " + answer, + NewRunConfig: func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "c", Stream: provider.StreamFnFromProvider(cp)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: agenttool.NewToolRegistry()}}, + } + }, + }) + } + a, b := mkChild("alpha"), mkChild("beta") + + // A single parent turn emitting two tool calls → both run in the same batch. + twoCall := fauxTurn{ + provider.StreamStartEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}}, + provider.StreamToolCallEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{ + agentcore.NewToolCallContent("c1", "agent-alpha", json.RawMessage(`{"prompt":"go"}`)), + agentcore.NewToolCallContent("c2", "agent-beta", json.RawMessage(`{"prompt":"go"}`)), + }}}, + provider.StreamDoneEvent{Message: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonToolUse, Content: agentcore.ContentList{ + agentcore.NewToolCallContent("c1", "agent-alpha", json.RawMessage(`{"prompt":"go"}`)), + agentcore.NewToolCallContent("c2", "agent-beta", json.RawMessage(`{"prompt":"go"}`)), + }}}, + } + parent := &fauxProvider{ + name: "faux-parent", + models: []provider.Model{{Provider: "p", ID: "p"}}, + turns: []fauxTurn{twoCall, textTurn("done")}, + } + cfg := newFauxRunCfg(parent, a, b) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("delegate both")}}, + }} + + _, msgs := collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + + got := map[string]string{} + for _, m := range msgs { + if tr, ok := m.(agentcore.ToolResultMessage); ok { + got[tr.ToolCallID] = textContentOf(tr.Content) + } + } + if got["c1"] != "alpha" || got["c2"] != "beta" { + t.Errorf("concurrent sub-agent results = %v, want c1=alpha c2=beta", got) + } +} + +// TestSubAgentEmptyPromptErrors verifies a sub-agent invoked with no prompt +// fails cleanly rather than spawning an empty child. +func TestSubAgentEmptyPromptErrors(t *testing.T) { + sub := NewSubAgentTool(SubAgentSpec{ + Name: "x", + Description: "x", + NewRunConfig: func() RunConfig { return RunConfig{} }, + }) + if _, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":""}`), nil); err == nil { + t.Error("empty prompt must error") + } +} + +// TestSubAgentFailedChildErrors verifies a child whose final turn stopped on +// error/aborted is surfaced to the parent as a tool error (not a silent +// success), so the parent model learns the delegation failed. +func TestSubAgentFailedChildErrors(t *testing.T) { + // A child turn that ends with StopReason=error carrying diagnostic text. + errTurn := func(text string) fauxTurn { + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + withText := partial + withText.Content = agentcore.ContentList{agentcore.NewTextContent(text)} + final := withText + final.StopReason = agentcore.StopReasonError + return fauxTurn{ + provider.StreamStartEvent{Partial: partial}, + provider.StreamTextEvent{Partial: withText}, + provider.StreamDoneEvent{Message: final}, + } + } + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{errTurn("provider blew up")}, + } + sub := NewSubAgentTool(SubAgentSpec{ + Name: "researcher", + Description: "delegate", + NewRunConfig: func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: provider.StreamFnFromProvider(child)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: agenttool.NewToolRegistry()}}, + } + }, + }) + _, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil) + if err == nil { + t.Fatal("a child that stopped on error must surface as a tool error") + } + if !strings.Contains(err.Error(), "provider blew up") { + t.Errorf("error should carry the child's diagnostic text, got %v", err) + } +} + +// --- Skills ----------------------------------------------------------------- + +// TestParseSkill verifies frontmatter + body parsing, including the name/body +// split and the required-description guard. +func TestParseSkill(t *testing.T) { + content := []byte("---\nname: summarize\ndescription: summarize a file\nallowed-tools:\n - read\n---\nYou summarize files.\nBe concise.\n") + sk, err := ParseSkill("summarize.md", content) + if err != nil { + t.Fatalf("ParseSkill: %v", err) + } + if sk.Frontmatter.Name != "summarize" { + t.Errorf("name = %q, want summarize", sk.Frontmatter.Name) + } + if sk.Frontmatter.Description != "summarize a file" { + t.Errorf("description = %q", sk.Frontmatter.Description) + } + if len(sk.Frontmatter.AllowedTools) != 1 || sk.Frontmatter.AllowedTools[0] != "read" { + t.Errorf("allowed-tools = %v, want [read]", sk.Frontmatter.AllowedTools) + } + if !strings.Contains(sk.Body, "You summarize files.") { + t.Errorf("body missing instructions: %q", sk.Body) + } +} + +// TestParseSkillDefaultsNameToFile verifies a skill without an explicit name +// defaults to its file base name. +func TestParseSkillDefaultsNameToFile(t *testing.T) { + sk, err := ParseSkill("/skills/deploy.md", []byte("---\ndescription: deploys\n---\nbody")) + if err != nil { + t.Fatalf("ParseSkill: %v", err) + } + if sk.Frontmatter.Name != "deploy" { + t.Errorf("name defaulted to %q, want deploy", sk.Frontmatter.Name) + } +} + +// TestParseSkillRequiresDescription verifies the description guard. +func TestParseSkillRequiresDescription(t *testing.T) { + if _, err := ParseSkill("x.md", []byte("---\nname: x\n---\nbody")); err == nil { + t.Error("skill without description must error") + } + if _, err := ParseSkill("x.md", []byte("no frontmatter here")); err == nil { + t.Error("skill without frontmatter must error") + } +} + +// TestParseSkillDisableModelInvocation verifies the disable-model-invocation +// frontmatter key parses into the three expected states. +func TestParseSkillDisableModelInvocation(t *testing.T) { + cases := []struct { + name string + yaml string + want bool + }{ + {"absent defaults false", "---\nname: a\ndescription: d\n---\nbody", false}, + {"explicit true", "---\nname: a\ndescription: d\ndisable-model-invocation: true\n---\nbody", true}, + {"explicit false", "---\nname: a\ndescription: d\ndisable-model-invocation: false\n---\nbody", false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + sk, err := ParseSkill("a.md", []byte(tc.yaml)) + if err != nil { + t.Fatalf("ParseSkill: %v", err) + } + if sk.Frontmatter.DisableModelInvocation != tc.want { + t.Errorf("DisableModelInvocation = %v, want %v", sk.Frontmatter.DisableModelInvocation, tc.want) + } + }) + } +} + +// TestValidateSkillName verifies the Agent Skills name rules: lowercase +// a-z/0-9/hyphen only, at most 64 chars, no leading/trailing/consecutive +// hyphens. +func TestValidateSkillName(t *testing.T) { + valid := []string{"weather", "chao-go-sync", "a", "a1-b2"} + for _, n := range valid { + if err := validateSkillName(n); err != nil { + t.Errorf("validateSkillName(%q) = %v, want nil", n, err) + } + } + invalid := []string{ + "Weather", // uppercase + "my_skill", // underscore + "has space", // space + "plugin:skill", // colon + "-lead", // leading hyphen + "trail-", // trailing hyphen + "double--hyphen", // consecutive hyphens + strings.Repeat("a", 65), // too long + } + for _, n := range invalid { + if err := validateSkillName(n); err == nil { + t.Errorf("validateSkillName(%q) = nil, want error", n) + } + } +} + +// TestValidateSkillDescription verifies description is required and bounded. +func TestValidateSkillDescription(t *testing.T) { + if err := validateSkillDescription("does a thing"); err != nil { + t.Errorf("valid description rejected: %v", err) + } + if err := validateSkillDescription(" "); err == nil { + t.Error("blank description must error") + } + if err := validateSkillDescription(strings.Repeat("x", 1025)); err == nil { + t.Error("over-long description must error") + } + if err := validateSkillDescription(strings.Repeat("x", 1024)); err != nil { + t.Errorf("1024-char description rejected: %v", err) + } +} + +// TestParseSkillRejectsInvalidName verifies an invalid name (including one +// derived from the file base name) makes ParseSkill fail, so LoadSkillsDir +// skips it and accumulates the reason rather than surfacing a bad skill. +func TestParseSkillRejectsInvalidName(t *testing.T) { + if _, err := ParseSkill("x.md", []byte("---\nname: Bad_Name\ndescription: d\n---\nbody")); err == nil { + t.Error("invalid explicit name must error") + } + if _, err := ParseSkill("/skills/My_Skill.md", []byte("---\ndescription: d\n---\nbody")); err == nil { + t.Error("invalid file-derived name must error") + } +} + +// TestFormatSkillsForPrompt verifies the block lists visible +// skills with name/description/location and excludes disabled ones. +func TestFormatSkillsForPrompt(t *testing.T) { + skills := []*Skill{ + {Frontmatter: SkillFrontmatter{Name: "weather", Description: "get weather"}, Path: "/skills/weather.md"}, + {Frontmatter: SkillFrontmatter{Name: "secret", Description: "hidden", DisableModelInvocation: true}, Path: "/skills/secret.md"}, + } + out := FormatSkillsForPrompt(skills) + if !strings.Contains(out, "") || !strings.Contains(out, "") { + t.Fatalf("missing block wrapper:\n%s", out) + } + if !strings.Contains(out, "weather") { + t.Errorf("visible skill name missing:\n%s", out) + } + if !strings.Contains(out, "get weather") { + t.Errorf("visible skill description missing:\n%s", out) + } + if !strings.Contains(out, "/skills/weather.md") { + t.Errorf("visible skill location missing:\n%s", out) + } + if strings.Contains(out, "secret") { + t.Errorf("disabled skill must be excluded:\n%s", out) + } + if !strings.Contains(out, "Use the read tool to load a skill's file") { + t.Errorf("guidance preamble missing:\n%s", out) + } +} + +// TestFormatSkillsForPromptEmpty verifies an empty or all-disabled list yields +// the empty string so callers can append unconditionally. +func TestFormatSkillsForPromptEmpty(t *testing.T) { + if got := FormatSkillsForPrompt(nil); got != "" { + t.Errorf("nil skills = %q, want empty", got) + } + disabled := []*Skill{{Frontmatter: SkillFrontmatter{Name: "x", Description: "d", DisableModelInvocation: true}, Path: "x.md"}} + if got := FormatSkillsForPrompt(disabled); got != "" { + t.Errorf("all-disabled skills = %q, want empty", got) + } +} + +// TestFormatSkillsForPromptEscapesXML verifies XML special characters in name +// and description are escaped. +func TestFormatSkillsForPromptEscapesXML(t *testing.T) { + skills := []*Skill{ + {Frontmatter: SkillFrontmatter{Name: "a", Description: `x & y < z > "q" 'r'`}, Path: "/s/a.md"}, + } + out := FormatSkillsForPrompt(skills) + if !strings.Contains(out, "x & y < z > "q" 'r'") { + t.Errorf("XML not escaped:\n%s", out) + } + if strings.Contains(out, "& y") || strings.Contains(out, "< z") { + t.Errorf("raw XML chars leaked:\n%s", out) + } +} + +// TestFormatSkillsForPromptAbsoluteLocation verifies a relative skill path is +// rendered as an absolute location so the model can read it from any cwd. +func TestFormatSkillsForPromptAbsoluteLocation(t *testing.T) { + skills := []*Skill{ + {Frontmatter: SkillFrontmatter{Name: "rel", Description: "d"}, Path: "sub/rel.md"}, + } + out := FormatSkillsForPrompt(skills) + if !strings.Contains(out, ""+string(filepath.Separator)) && !strings.Contains(out, "/") { + t.Errorf("location should be absolute:\n%s", out) + } +} + +// TestLoadSkillsDir verifies loading flat *.md skills and nested /SKILL.md, +// sorted by name; a missing dir is not an error. +func TestLoadSkillsDir(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "beta.md"), []byte("---\ndescription: beta skill\n---\nbeta body"), 0o644); err != nil { + t.Fatal(err) + } + nested := filepath.Join(dir, "alpha") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(nested, "SKILL.md"), []byte("---\nname: alpha\ndescription: alpha skill\n---\nalpha body"), 0o644); err != nil { + t.Fatal(err) + } + + skills, err := LoadSkillsDir(dir) + if err != nil { + t.Fatalf("LoadSkillsDir: %v", err) + } + if len(skills) != 2 { + t.Fatalf("loaded %d skills, want 2", len(skills)) + } + if skills[0].Frontmatter.Name != "alpha" || skills[1].Frontmatter.Name != "beta" { + t.Errorf("skills not sorted by name: %q, %q", skills[0].Frontmatter.Name, skills[1].Frontmatter.Name) + } + + // Missing directory → no skills, no error. + empty, err := LoadSkillsDir(filepath.Join(dir, "does-not-exist")) + if err != nil || empty != nil { + t.Errorf("missing dir should yield (nil, nil), got (%v, %v)", empty, err) + } +} + +// TestSkillSubAgentSpec verifies a skill materializes as a sub-agent whose +// system prompt is the body, whose tools are filtered by allowed-tools, and +// whose description is surfaced. +func TestSkillSubAgentSpec(t *testing.T) { + sk := &Skill{ + Frontmatter: SkillFrontmatter{Name: "reader", Description: "reads", AllowedTools: []string{"read"}}, + Body: "you read files", + } + tools := []agentcore.AgentTool{execTool{name: "read"}, execTool{name: "write"}, execTool{name: "bash"}} + var gotTools []agentcore.AgentTool + spec := sk.SubAgentSpec(tools, func(resolved []agentcore.AgentTool) RunConfig { + gotTools = resolved + return RunConfig{} + }) + if spec.Name != "reader" || spec.Description != "reads" { + t.Errorf("spec identity = %q/%q", spec.Name, spec.Description) + } + if spec.SystemPrompt != "you read files" { + t.Errorf("spec system prompt = %q", spec.SystemPrompt) + } + if len(spec.Tools) != 1 || spec.Tools[0].Name() != "read" { + t.Errorf("allowed-tools filter failed, spec tools = %v", spec.Tools) + } + // NewRunConfig passes the resolved (filtered) tool set to the factory. + spec.NewRunConfig() + if len(gotTools) != 1 || gotTools[0].Name() != "read" { + t.Errorf("factory received %v, want [read]", gotTools) + } +} + +// --- Slash-commands --------------------------------------------------------- + +// TestSlashBuiltinWinsOverUser is acceptance-critical: the conflict priority +// rule keeps a built-in over a same-named user command, recording the shadow. +func TestSlashBuiltinWinsOverUser(t *testing.T) { + // Register a built-in under a unique name to avoid cross-test pollution. + name := "compact-test-builtin" + if _, exists := builtinCommands[name]; !exists { + RegisterBuiltin(SlashCommand{Name: name, Description: "builtin", Expand: func(string) string { return "BUILTIN" }}) + } + r := NewSlashRegistry() + r.AddUser(SlashCommand{Name: name, Expand: func(string) string { return "USER" }}) + + cmd, ok := r.Lookup(name) + if !ok { + t.Fatalf("command %q not found", name) + } + if cmd.Source != SourceBuiltin { + t.Errorf("built-in must win, got source %v", cmd.Source) + } + if got := cmd.Expand(""); got != "BUILTIN" { + t.Errorf("expanded %q, want BUILTIN (built-in wins)", got) + } + shadowed := r.Shadowed() + if len(shadowed) != 1 || shadowed[0].Name != name { + t.Errorf("shadowed = %v, want [%s]", shadowed, name) + } +} + +// TestSlashResolve verifies "/name args" parsing, non-command passthrough, and +// the unknown-command error. +func TestSlashResolve(t *testing.T) { + r := NewSlashRegistry() + r.AddUser(SlashCommand{Name: "greet", Expand: func(args string) string { return "hello " + args }}) + + // Slash command with args → expanded. + prompt, handled, err := r.Resolve("/greet world") + if err != nil || !handled || prompt != "hello world" { + t.Errorf("Resolve(/greet world) = (%q, %v, %v)", prompt, handled, err) + } + // Non-command passthrough. + prompt, handled, err = r.Resolve("just a normal prompt") + if err != nil || handled || prompt != "just a normal prompt" { + t.Errorf("non-command passthrough failed: (%q, %v, %v)", prompt, handled, err) + } + // Unknown command errors. + if _, _, err := r.Resolve("/nope"); err == nil { + t.Error("unknown command must error") + } +} + +// TestSlashActionCommand verifies an action command runs its side effect via +// ResolveOutcome and reports SlashAction with its status message (no prompt), +// while a prompt command reports SlashPrompt with expanded text. +func TestSlashActionCommand(t *testing.T) { + r := NewSlashRegistry() + var ran string + r.AddBuiltin(SlashCommand{ + Name: "model", + Description: "switch model", + Action: func(args string) string { ran = args; return "switched to " + args }, + }) + r.AddUser(SlashCommand{Name: "greet", Expand: func(args string) string { return "hello " + args }}) + + // Action command: runs the side effect and returns a status message. + out, err := r.ResolveOutcome("/model gpt-5") + if err != nil { + t.Fatalf("ResolveOutcome(/model) error: %v", err) + } + if !out.Handled || out.Kind != SlashAction { + t.Errorf("action command: got handled=%v kind=%v, want true/SlashAction", out.Handled, out.Kind) + } + if ran != "gpt-5" { + t.Errorf("action side effect not run with args: got %q", ran) + } + if out.Message != "switched to gpt-5" || out.Prompt != "" { + t.Errorf("action outcome = {msg:%q prompt:%q}, want status message and empty prompt", out.Message, out.Prompt) + } + + // Prompt command: expands, no action. + out, err = r.ResolveOutcome("/greet world") + if err != nil { + t.Fatalf("ResolveOutcome(/greet) error: %v", err) + } + if !out.Handled || out.Kind != SlashPrompt || out.Prompt != "hello world" { + t.Errorf("prompt outcome = {kind:%v prompt:%q}, want SlashPrompt/hello world", out.Kind, out.Prompt) + } + + // Non-command passthrough. + out, err = r.ResolveOutcome("plain text") + if err != nil || out.Handled || out.Prompt != "plain text" { + t.Errorf("passthrough = {%v %q %v}, want unhandled verbatim", out.Handled, out.Prompt, err) + } +} + +// TestAddBuiltinDuplicatePanics verifies AddBuiltin rejects a duplicate +// built-in name (a programming error), matching RegisterBuiltin semantics. +func TestAddBuiltinDuplicatePanics(t *testing.T) { + r := NewSlashRegistry() + r.AddBuiltin(SlashCommand{Name: "dup", Action: func(string) string { return "" }}) + defer func() { + if recover() == nil { + t.Error("duplicate AddBuiltin must panic") + } + }() + r.AddBuiltin(SlashCommand{Name: "dup", Action: func(string) string { return "" }}) +} + +// TestAddBuiltinWinsOverUser verifies an instance built-in (AddBuiltin) shadows +// a same-named user command, just like a globally registered built-in. +func TestAddBuiltinWinsOverUser(t *testing.T) { + r := NewSlashRegistry() + r.AddBuiltin(SlashCommand{Name: "x", Action: func(string) string { return "builtin" }}) + r.AddUser(SlashCommand{Name: "x", Expand: func(string) string { return "user" }}) + cmd, ok := r.Lookup("x") + if !ok || cmd.Source != SourceBuiltin { + t.Errorf("instance built-in must win, got ok=%v source=%v", ok, cmd.Source) + } + if len(r.Shadowed()) != 1 || r.Shadowed()[0].Name != "x" { + t.Errorf("shadowed = %v, want [x]", r.Shadowed()) + } +} + +// TestParseUserCommand verifies $ARGUMENTS substitution, frontmatter description, +// and the append fallback when no placeholder is present. +func TestParseUserCommand(t *testing.T) { + // With frontmatter + placeholder. + cmd, err := ParseUserCommand("review", []byte("---\ndescription: review code\n---\nReview this: $ARGUMENTS please")) + if err != nil { + t.Fatalf("ParseUserCommand: %v", err) + } + if cmd.Description != "review code" { + t.Errorf("description = %q", cmd.Description) + } + if got := cmd.Expand("main.go"); got != "Review this: main.go please" { + t.Errorf("expand = %q", got) + } + // No placeholder: args appended. + bare, _ := ParseUserCommand("note", []byte("Take a note")) + if got := bare.Expand("buy milk"); got != "Take a note\n\nbuy milk" { + t.Errorf("bare expand = %q", got) + } + if got := bare.Expand(""); got != "Take a note" { + t.Errorf("bare expand no args = %q", got) + } +} + +// TestLoadUserCommandsDir verifies loading *.md command templates from a dir, +// sorted by name; a missing dir is not an error. +func TestLoadUserCommandsDir(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "deploy.md"), []byte("Deploy $ARGUMENTS now"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "test.md"), []byte("---\ndescription: run tests\n---\nRun tests"), 0o644); err != nil { + t.Fatal(err) + } + cmds, err := LoadUserCommandsDir(dir) + if err != nil { + t.Fatalf("LoadUserCommandsDir: %v", err) + } + if len(cmds) != 2 || cmds[0].Name != "deploy" || cmds[1].Name != "test" { + t.Fatalf("loaded %d commands (want deploy,test sorted): %+v", len(cmds), cmds) + } + if got := cmds[0].Expand("prod"); got != "Deploy prod now" { + t.Errorf("deploy expand = %q", got) + } + + empty, err := LoadUserCommandsDir(filepath.Join(dir, "missing")) + if err != nil || empty != nil { + t.Errorf("missing dir should yield (nil, nil), got (%v, %v)", empty, err) + } +} diff --git a/pigo/internal/runtime/progress_test.go b/pigo/internal/runtime/progress_test.go new file mode 100644 index 0000000..4bc66a0 --- /dev/null +++ b/pigo/internal/runtime/progress_test.go @@ -0,0 +1,149 @@ +package runtime + +// Tests for sub-agent progress reporting (US-005, #455): a dispatched task's +// child tool-execution / turn boundaries are translated into +// SubAgentProgressEvent and surfaced through the run-level emitter the parent +// loop injects into ctx (WithProgressEmitter). The parent tool-call id must ride +// on the event, the activity must map from the child event, and a nil emitter +// (the tool called outside a loop, e.g. a direct unit test) must be a silent +// no-op rather than a panic. The child loop is driven through the faux provider +// seam; only the provider boundary is faked. + +import ( + "context" + "encoding/json" + "sync" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/provider" +) + +// TestActivityOf pins the child-event → activity mapping (D-8 / §5.3): tool +// starts map to their display verb, a turn start maps to "Thinking", and +// everything else maps to "" (no emission). +func TestActivityOf(t *testing.T) { + cases := []struct { + ev agentcore.AgentEvent + want string + }{ + {agentcore.ToolExecutionStartEvent{ToolName: "read"}, "Reading"}, + {agentcore.ToolExecutionStartEvent{ToolName: "edit"}, "Editing"}, + {agentcore.ToolExecutionStartEvent{ToolName: "write"}, "Editing"}, + {agentcore.ToolExecutionStartEvent{ToolName: "bash"}, "Running bash"}, + {agentcore.ToolExecutionStartEvent{ToolName: "grep"}, "Searching"}, + {agentcore.ToolExecutionStartEvent{ToolName: "find"}, "Searching"}, + {agentcore.ToolExecutionStartEvent{ToolName: "ls"}, "Searching"}, + {agentcore.ToolExecutionStartEvent{ToolName: "webfetch"}, "Fetching"}, + {agentcore.ToolExecutionStartEvent{ToolName: "todo"}, ""}, + {agentcore.TurnStartEvent{}, "Thinking"}, + {agentcore.ToolExecutionEndEvent{ToolName: "read"}, ""}, + {agentcore.MessageStartEvent{}, ""}, + } + for _, c := range cases { + if got := activityOf(c.ev); got != c.want { + t.Errorf("activityOf(%T{%v}) = %q, want %q", c.ev, c.ev, got, c.want) + } + } +} + +// TestTaskEmitsSubAgentProgress verifies a child that executes a tool triggers a +// SubAgentProgressEvent carrying the parent task's tool-call id and the mapped +// activity ("Reading" for a child "read" tool), plus the "Thinking" boundary at +// each turn start. +func TestTaskEmitsSubAgentProgress(t *testing.T) { + // A child tool named "read" so its ToolExecutionStart maps to "Reading". + readTool := execTool{ + name: "read", + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("file body")}}, nil + }, + } + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{toolCallTurn("t1", "read", `{}`), textTurn("child final report")}, + } + factory := func() RunConfig { + reg := agenttool.NewToolRegistry() + _ = reg.Register(readTool) + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: provider.StreamFnFromProvider(child)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}}, + } + } + tool := NewTaskTool(factory, nil) + + var mu sync.Mutex + var progress []agentcore.SubAgentProgressEvent + emit := func(ctx context.Context, ev agentcore.AgentEvent) error { + if p, ok := ev.(agentcore.SubAgentProgressEvent); ok { + mu.Lock() + progress = append(progress, p) + mu.Unlock() + } + return nil + } + ctx := agentcore.WithProgressEmitter(context.Background(), emit) + + const parentID = "parent-call-id" + res, err := tool.Execute(ctx, parentID, json.RawMessage(`{"description":"read the file","prompt":"do the work"}`), nil) + if err != nil { + t.Fatalf("Execute err = %v", err) + } + if got := agentcore.ContentToText(res.Content); got != "child final report" { + t.Errorf("task result = %q, want 'child final report'", got) + } + + mu.Lock() + defer mu.Unlock() + if len(progress) == 0 { + t.Fatal("expected at least one SubAgentProgressEvent, got none") + } + var sawReading bool + for _, p := range progress { + if p.ToolCallID != parentID { + t.Errorf("progress ToolCallID = %q, want %q", p.ToolCallID, parentID) + } + if p.Description != "read the file" { + t.Errorf("progress Description = %q, want 'read the file'", p.Description) + } + if p.Activity == "" { + t.Errorf("progress emitted with empty Activity (should be skipped)") + } + if p.Activity == "Reading" { + sawReading = true + } + } + if !sawReading { + t.Errorf("expected a 'Reading' activity from the child read tool, got %+v", progress) + } +} + +// TestTaskNilEmitterNoPanic verifies that when no progress emitter is present in +// ctx (the tool called outside a loop, as in direct unit tests) the child still +// runs, returns its text, and emits no progress — without panicking. +func TestTaskNilEmitterNoPanic(t *testing.T) { + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{textTurn("child final report")}, + } + factory := func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: provider.StreamFnFromProvider(child)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: agenttool.NewToolRegistry()}}, + } + } + tool := NewTaskTool(factory, nil) + + // Plain ctx: no WithProgressEmitter, so ProgressEmitterFromContext is nil. + res, err := tool.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil) + if err != nil { + t.Fatalf("Execute err = %v", err) + } + if got := agentcore.ContentToText(res.Content); got != "child final report" { + t.Errorf("task result = %q, want 'child final report'", got) + } +} diff --git a/pigo/internal/runtime/prompt.go b/pigo/internal/runtime/prompt.go new file mode 100644 index 0000000..6082b98 --- /dev/null +++ b/pigo/internal/runtime/prompt.go @@ -0,0 +1,203 @@ +// This file implements system-prompt assembly (US-021, #40), the pigo port of +// pi's prompt construction. A run's system prompt is built from three layers, +// in order: +// +// 1. a base instruction (who the agent is and how it should behave), +// 2. an environment block (working directory, OS/arch, current date), and +// 3. every AGENTS.md found on the path from a root directory down to the +// working directory, concatenated general-to-specific. +// +// The AGENTS.md ordering mirrors zero/pi's monorepo behavior: a repo-root +// AGENTS.md states broad conventions, and a nested package's AGENTS.md refines +// them, so the more specific file appears later and takes precedence in the +// model's reading. +package runtime + +import ( + "fmt" + "os" + "path/filepath" + "runtime" + "strings" + "time" +) + +// agentsFileName is the per-directory instruction file injected into the system +// prompt, general-to-specific from the root down to the working directory. +const agentsFileName = "AGENTS.md" + +// PromptConfig configures system-prompt assembly. The zero value is usable: it +// produces the base instruction plus an environment block for the process +// working directory, with no AGENTS.md injection. +type PromptConfig struct { + // BaseInstruction is the leading text of the system prompt. When empty, + // DefaultBaseInstruction is used. + BaseInstruction string + // WorkingDir is the directory the run operates in. When empty, the process + // working directory (os.Getwd) is used. It anchors both the environment + // block and the lower bound of the AGENTS.md walk. + WorkingDir string + // Root bounds the AGENTS.md walk at its top. AGENTS.md files are injected for + // every directory from Root down to WorkingDir, inclusive. When empty, only + // WorkingDir's own AGENTS.md (if any) is considered — no ancestor walk. + Root string + // AppendInstructions are appended verbatim to the end of the assembled + // prompt, in order, each preceded by a blank line. This is the sink for + // --append-system-prompt (mirrors pi): extra guidance layered after the base + // instruction, environment block, and AGENTS.md. Empty entries are skipped. + AppendInstructions []string + // Now supplies the timestamp for the environment block. When nil, time.Now + // is used. Injected for deterministic tests. + Now func() time.Time + // ReadFile reads a file's contents. When nil, os.ReadFile is used. Injected + // for tests so AGENTS.md layout can be faked without touching disk. + ReadFile func(path string) ([]byte, error) + // Skills are the model-invocable skills to advertise in an + // block at the end of the prompt (mirrors pi's progressive disclosure). Only their + // name/description/location are injected; the model loads a skill's body with + // the read tool on demand. Skills flagged disable-model-invocation are + // filtered out by FormatSkillsForPrompt. Empty means no skills block. + Skills []*Skill + // ReadToolAvailable signals whether the read tool is in the current tool set. + // Skills are advertised only when it is true, since the model needs the read + // tool to load a skill's body; otherwise the block is omitted entirely. + ReadToolAvailable bool +} + +// DefaultBaseInstruction is the leading system-prompt text used when +// PromptConfig.BaseInstruction is empty. +const DefaultBaseInstruction = "You are pigo, a helpful coding agent. " + + "Use the available tools to inspect files and accomplish the user's request precisely and concisely.\n\n" + + todoGuide + "\n\n" + + taskGuide + +// todoGuide instructs the model on how to drive the todo tool. It is appended to +// the default base instruction so multi-step work is planned and its progress is +// made visible to the user (US-011). +const todoGuide = "When a task has multiple steps or is non-trivial, use the todo tool to plan " + + "and track your work. Submit the entire task list each call (it replaces the previous " + + "list); each item has a content string and a status of pending, in_progress, or completed. " + + "Keep exactly one item in_progress at a time, and mark an item completed as soon as it is " + + "done before starting the next. Skip the todo tool for trivial single-step requests." + +// taskGuide instructs the model on how to use the generic `task` tool (US-008, +// #458). The task tool dispatches an independent sub-agent that runs its own +// agent loop with a fresh context and returns its final report, so it is the +// mechanism for delegation and fan-out. The key affordance advertised here is +// that emitting MULTIPLE task calls in a single assistant message runs those +// sub-agents in parallel, letting skills like /graph achieve real concurrency. +const taskGuide = "When work splits into independent subtasks, delegate them with the task tool: each " + + "task call dispatches an independent sub-agent that completes its subtask on a fresh context and " + + "returns its final report. To fan out, emit MULTIPLE task calls in a single message — they run in " + + "parallel. Give each a complete, self-contained prompt, since a sub-agent shares none of this " + + "conversation's context. Do the work directly for a single, sequential, or trivial task." + +// BuildSystemPrompt assembles the full system prompt from cfg: base instruction, +// environment block, then AGENTS.md files ordered general-to-specific from Root +// down to WorkingDir. Missing AGENTS.md files are skipped silently; only a +// present-but-unreadable file (a real I/O error other than not-exist) is +// reported. +func BuildSystemPrompt(cfg PromptConfig) (string, error) { + base := cfg.BaseInstruction + if base == "" { + base = DefaultBaseInstruction + } + wd := cfg.WorkingDir + if wd == "" { + if cwd, err := os.Getwd(); err == nil { + wd = cwd + } + } + now := cfg.Now + if now == nil { + now = time.Now + } + readFile := cfg.ReadFile + if readFile == nil { + readFile = os.ReadFile + } + + var b strings.Builder + b.WriteString(base) + + b.WriteString("\n\nEnvironment:\n") + fmt.Fprintf(&b, "- Working directory: %s\n", wd) + fmt.Fprintf(&b, "- OS: %s/%s\n", runtime.GOOS, runtime.GOARCH) + fmt.Fprintf(&b, "- Date: %s", now().Format("2006-01-02")) + + dirs := agentsDirChain(cfg.Root, wd) + for _, dir := range dirs { + path := filepath.Join(dir, agentsFileName) + data, err := readFile(path) + if err != nil { + if os.IsNotExist(err) { + continue + } + return "", fmt.Errorf("read %s: %w", path, err) + } + content := strings.TrimSpace(string(data)) + if content == "" { + continue + } + fmt.Fprintf(&b, "\n\n# Project instructions (%s)\n%s", path, content) + } + + // Appended instructions (--append-system-prompt) come last so they layer on + // top of the base instruction, environment, and AGENTS.md. Each is separated + // by a blank line; empty entries are skipped. + for _, extra := range cfg.AppendInstructions { + extra = strings.TrimSpace(extra) + if extra == "" { + continue + } + b.WriteString("\n\n") + b.WriteString(extra) + } + + // Advertise model-invocable skills last (progressive disclosure), but only + // when the read tool is available — the model needs it to load a skill's + // body. FormatSkillsForPrompt returns "" when no visible skill remains, so + // this leaves a skill-free prompt byte-for-byte unchanged. + if cfg.ReadToolAvailable { + b.WriteString(FormatSkillsForPrompt(cfg.Skills)) + } + + return b.String(), nil +} + +// agentsDirChain returns the directories whose AGENTS.md should be injected, in +// general-to-specific order (root first, working directory last). When root is +// empty or is not an ancestor of wd, only wd is returned. When wd is empty, the +// chain is empty. +func agentsDirChain(root, wd string) []string { + if wd == "" { + return nil + } + wd = filepath.Clean(wd) + if root == "" { + return []string{wd} + } + root = filepath.Clean(root) + + // Walk up from wd to root, collecting each directory, then reverse so the + // root comes first (general → specific). If root is never reached, wd is not + // under root, so fall back to wd alone. + var up []string + cur := wd + for { + up = append(up, cur) + if cur == root { + // Reverse in place: root-first. + for i, j := 0, len(up)-1; i < j; i, j = i+1, j-1 { + up[i], up[j] = up[j], up[i] + } + return up + } + parent := filepath.Dir(cur) + if parent == cur { + // Reached filesystem root without hitting `root`: wd not under root. + return []string{wd} + } + cur = parent + } +} diff --git a/pigo/internal/runtime/prompt_test.go b/pigo/internal/runtime/prompt_test.go new file mode 100644 index 0000000..bff3c91 --- /dev/null +++ b/pigo/internal/runtime/prompt_test.go @@ -0,0 +1,315 @@ +package runtime + +// Tests for system-prompt assembly (US-021, #40): the base instruction, the +// environment block, and — the acceptance-critical part — the general-to- +// specific ordering of AGENTS.md injection from a root directory down to the +// working directory. AGENTS.md layout is faked via PromptConfig.ReadFile so the +// ordering is asserted without touching disk. + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// fixedTime is a deterministic clock for the environment block. +func fixedTime() time.Time { return time.Date(2026, 7, 10, 12, 0, 0, 0, time.UTC) } + +// TestBuildSystemPromptBaseAndEnv verifies the base instruction and environment +// block (cwd, OS, date) are present, with no AGENTS.md when none exist. +func TestBuildSystemPromptBaseAndEnv(t *testing.T) { + got, err := BuildSystemPrompt(PromptConfig{ + WorkingDir: "/work/proj", + Now: fixedTime, + ReadFile: func(string) ([]byte, error) { return nil, os.ErrNotExist }, + }) + if err != nil { + t.Fatalf("BuildSystemPrompt: %v", err) + } + if !strings.HasPrefix(got, DefaultBaseInstruction) { + t.Errorf("prompt should start with the default base instruction, got:\n%s", got) + } + if !strings.Contains(got, "Working directory: /work/proj") { + t.Errorf("environment block missing working directory:\n%s", got) + } + if !strings.Contains(got, "Date: 2026-07-10") { + t.Errorf("environment block missing date:\n%s", got) + } + if strings.Contains(got, "Project instructions") { + t.Errorf("no AGENTS.md exists, but prompt injected one:\n%s", got) + } +} + +// TestBuildSystemPromptAdvertisesTaskFanout verifies the base instruction tells +// the model about the generic task tool: that it dispatches an independent +// sub-agent (delegation) and that multiple task calls in one message run in +// parallel (fan-out, US-008/#458). +func TestBuildSystemPromptAdvertisesTaskFanout(t *testing.T) { + got, err := BuildSystemPrompt(PromptConfig{ + WorkingDir: "/work/proj", + Now: fixedTime, + ReadFile: func(string) ([]byte, error) { return nil, os.ErrNotExist }, + }) + if err != nil { + t.Fatalf("BuildSystemPrompt: %v", err) + } + lower := strings.ToLower(got) + if !strings.Contains(lower, "task tool") { + t.Errorf("prompt should advertise the task tool:\n%s", got) + } + if !strings.Contains(lower, "sub-agent") || !strings.Contains(lower, "independent") { + t.Errorf("prompt should describe the task tool as an independent sub-agent (delegation):\n%s", got) + } + if !strings.Contains(lower, "parallel") { + t.Errorf("prompt should state that multiple task calls run in parallel (fan-out):\n%s", got) + } +} + +// TestBuildSystemPromptAGENTSOrdering is the acceptance-critical test: with an +// AGENTS.md at the root and at a nested working directory, the root's content +// must appear BEFORE the nested one (general → specific). +func TestBuildSystemPromptAGENTSOrdering(t *testing.T) { + root := filepath.Clean("/repo") + mid := filepath.Join(root, "services") + wd := filepath.Join(mid, "api") + + files := map[string]string{ + filepath.Join(root, agentsFileName): "ROOT CONVENTIONS", + filepath.Join(mid, agentsFileName): "SERVICES CONVENTIONS", + filepath.Join(wd, agentsFileName): "API CONVENTIONS", + } + got, err := BuildSystemPrompt(PromptConfig{ + WorkingDir: wd, + Root: root, + Now: fixedTime, + ReadFile: func(path string) ([]byte, error) { + if c, ok := files[path]; ok { + return []byte(c), nil + } + return nil, os.ErrNotExist + }, + }) + if err != nil { + t.Fatalf("BuildSystemPrompt: %v", err) + } + + iRoot := strings.Index(got, "ROOT CONVENTIONS") + iMid := strings.Index(got, "SERVICES CONVENTIONS") + iAPI := strings.Index(got, "API CONVENTIONS") + if iRoot < 0 || iMid < 0 || iAPI < 0 { + t.Fatalf("all three AGENTS.md must be injected, got:\n%s", got) + } + if !(iRoot < iMid && iMid < iAPI) { + t.Errorf("AGENTS.md must be ordered general→specific (root= iAPI { + t.Errorf("present AGENTS.md must stay ordered root block is +// appended after the base/env/append layers when the read tool is available. +func TestBuildSystemPromptInjectsSkills(t *testing.T) { + skills := []*Skill{ + {Frontmatter: SkillFrontmatter{Name: "weather", Description: "get weather"}, Path: "/skills/weather.md"}, + {Frontmatter: SkillFrontmatter{Name: "secret", Description: "hidden", DisableModelInvocation: true}, Path: "/skills/secret.md"}, + } + got, err := BuildSystemPrompt(PromptConfig{ + WorkingDir: "/work/proj", + Now: fixedTime, + ReadFile: func(string) ([]byte, error) { return nil, os.ErrNotExist }, + Skills: skills, + ReadToolAvailable: true, + }) + if err != nil { + t.Fatalf("BuildSystemPrompt: %v", err) + } + if !strings.Contains(got, "") || !strings.Contains(got, "weather") { + t.Errorf("skills block must be injected, got:\n%s", got) + } + if strings.Contains(got, "secret") { + t.Errorf("disable-model-invocation skill must not appear, got:\n%s", got) + } + iEnv := strings.Index(got, "Working directory") + iSkills := strings.Index(got, "") + if !(iEnv < iSkills) { + t.Errorf("skills block must come after env block, env=%d skills=%d", iEnv, iSkills) + } +} + +// TestBuildSystemPromptNoSkillsWithoutReadTool verifies skills are NOT injected +// when the read tool is unavailable, and that a skill-free prompt is unchanged. +func TestBuildSystemPromptNoSkillsWithoutReadTool(t *testing.T) { + skills := []*Skill{{Frontmatter: SkillFrontmatter{Name: "weather", Description: "d"}, Path: "/s/weather.md"}} + withTool, _ := BuildSystemPrompt(PromptConfig{WorkingDir: "/w", Now: fixedTime, ReadFile: func(string) ([]byte, error) { return nil, os.ErrNotExist }, Skills: skills, ReadToolAvailable: false}) + if strings.Contains(withTool, "available_skills") { + t.Errorf("no read tool → no skills block, got:\n%s", withTool) + } + // A prompt with no skills at all must equal one with read tool but empty list. + bare, _ := BuildSystemPrompt(PromptConfig{WorkingDir: "/w", Now: fixedTime, ReadFile: func(string) ([]byte, error) { return nil, os.ErrNotExist }}) + withReadNoSkills, _ := BuildSystemPrompt(PromptConfig{WorkingDir: "/w", Now: fixedTime, ReadFile: func(string) ([]byte, error) { return nil, os.ErrNotExist }, ReadToolAvailable: true}) + if bare != withReadNoSkills { + t.Errorf("empty skill list must not alter the prompt even with read tool:\n%q\nvs\n%q", bare, withReadNoSkills) + } +} + +// TestDisableModelInvocationCoexistence verifies the #305 coexistence contract: +// a skill with disable-model-invocation:true is STILL exposed as a /skill-name +// slash command (body expansion + $ARGUMENTS), while being EXCLUDED from the +// prompt injection. The two invocation paths are independent. +func TestDisableModelInvocationCoexistence(t *testing.T) { + disabled := &Skill{ + Frontmatter: SkillFrontmatter{Name: "secret", Description: "hidden", DisableModelInvocation: true}, + Path: "/skills/secret.md", + Body: "Do the secret thing with $ARGUMENTS.", + } + enabled := &Skill{ + Frontmatter: SkillFrontmatter{Name: "weather", Description: "get weather"}, + Path: "/skills/weather.md", + Body: "Report the weather.", + } + skills := []*Skill{disabled, enabled} + + // 1. The disabled skill must be excluded from the model-facing prompt block, + // while the enabled one appears. + block := FormatSkillsForPrompt(skills) + if strings.Contains(block, "secret") { + t.Errorf("disable-model-invocation skill must not appear in , got:\n%s", block) + } + if !strings.Contains(block, "weather") { + t.Errorf("model-invocable skill must appear in , got:\n%s", block) + } + + // 2. The disabled skill must still be invocable via its /skill-name command, + // with $ARGUMENTS substitution intact (behavior identical to an enabled one). + cmd := disabled.SlashCommand() + if cmd.Name != "secret" { + t.Errorf("disabled skill slash name = %q, want secret", cmd.Name) + } + if cmd.Expand == nil { + t.Fatal("disabled skill must expose a prompt command (Expand != nil)") + } + if got := cmd.Expand("now"); got != "Do the secret thing with now." { + t.Errorf("Expand(now) = %q, want $ARGUMENTS substituted", got) + } +} diff --git a/pigo/internal/runtime/rebuild.go b/pigo/internal/runtime/rebuild.go new file mode 100644 index 0000000..dab9b1f --- /dev/null +++ b/pigo/internal/runtime/rebuild.go @@ -0,0 +1,165 @@ +// Context rebuild for the "infinite context" feature (#482). Where auto- +// compaction (loop.go) collapses history *lossily* on the fly, a rebuild +// reconstructs the working context deterministically from a persisted +// checkpoint: everything before the checkpoint watermark is replaced by the +// distilled checkpoint summary, and everything at/after the watermark is kept +// verbatim. This is what the /rebuild command (REPL + TUI) invokes, and what +// the run loop (#481) will call on resume to reload a collapsed prefix instead +// of replaying the whole transcript. +// +// When no checkpoint exists yet there is nothing to reload, so a rebuild falls +// back to the ordinary lossy compaction path (the same compaction.Compact flow +// runCompaction drives) so /rebuild still shrinks an overgrown context. +// +// This file is deliberately side-effect free with respect to the loop: it never +// mutates the caller's AgentContext. It returns the rebuilt MessageList (plus a +// RebuildResult describing what happened and an equivalent CompactionEvent) so +// the CLI handlers — and, later, #481 — decide when and how to apply it. +package runtime + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/compaction" +) + +// RebuildResult describes the outcome of a context rebuild. Messages is the +// rebuilt list ready to replace the live context; the remaining fields mirror +// CompactionEvent so a front-end can report the same before/after summary as a +// compaction. +type RebuildResult struct { + // Messages is the rebuilt context: a single summary/checkpoint message + // followed by the retained recent tail. When NoOp is true it is the original + // list, unchanged. + Messages agentcore.MessageList + // FromCheckpoint is true when the boundary came from a persisted checkpoint; + // false when the no-checkpoint fallback ran a lossy compaction. + FromCheckpoint bool + // Watermark is the boundary index used: the checkpoint watermark, or the + // compaction cut point in the fallback path. + Watermark int + // SummarizedCount is how many leading messages were collapsed into the summary. + SummarizedCount int + // KeptCount is how many recent messages were preserved verbatim. + KeptCount int + // TokensBefore / TokensAfter are the estimated context tokens before and after + // the rebuild (equal when NoOp). + TokensBefore int + TokensAfter int + // NoOp is true when nothing changed: no checkpoint existed and there was + // nothing to compact (an empty summarization range). + NoOp bool +} + +// Event renders the rebuild as a CompactionEvent so consumers that already +// handle compaction reporting (the REPL/TUI event surfaces) can present a +// rebuild with the same shape. Reason is "rebuild". +func (r *RebuildResult) Event() agentcore.CompactionEvent { + return agentcore.CompactionEvent{ + Reason: "rebuild", + TokensBefore: r.TokensBefore, + TokensAfter: r.TokensAfter, + SummarizedCount: r.SummarizedCount, + KeptCount: r.KeptCount, + } +} + +// RebuildFromCheckpoint reconstructs the working context for sessionID. +// +// If a checkpoint exists under memoryRoot, the compression boundary is inserted +// at checkpoint.Watermark: messages before the watermark collapse to the +// checkpoint summary (rendered as a single compaction message), and messages +// at/after the watermark are preserved verbatim. The watermark is clamped to +// [0, len(msgs)] so a stale checkpoint recorded against a longer history (or one +// that has since been re-compacted) never slices out of range. +// +// If no checkpoint exists, it falls back to the lossy compaction path — the same +// compaction.Compact flow the loop's auto-compaction uses (runCompaction) — so +// /rebuild still shrinks the context. When there is nothing to compact the +// original list is returned with NoOp set. +// +// It performs no mutation of the caller's context and no checkpoint writes; the +// returned RebuildResult carries the rebuilt list for the caller to apply. When +// waitForCheckpoint is non-nil it is invoked before reading, so a caller that +// has an in-flight checkpoint write can block until it lands (kept as a callback +// so this stays decoupled from the loop's write path). +func RebuildFromCheckpoint( + ctx context.Context, + msgs agentcore.MessageList, + sessionID, memoryRoot string, + cfg *RunConfig, + waitForCheckpoint func(), +) (*RebuildResult, error) { + if waitForCheckpoint != nil { + waitForCheckpoint() + } + now := nowMillis() + tokensBefore := compaction.EstimateContextTokens(msgs).Tokens + + cp, ok, err := LoadCheckpoint(sessionID, memoryRoot) + if err != nil { + return nil, err + } + if ok { + return rebuildFromLoadedCheckpoint(msgs, cp, tokensBefore, now), nil + } + + // No checkpoint: fall back to the ordinary lossy compaction path. + res, err := runCompaction(ctx, msgs, cfg) + if err != nil { + return nil, err + } + if res == nil { + // Nothing to summarize (no valid cut point / empty range): leave as-is. + return &RebuildResult{ + Messages: msgs, + TokensBefore: tokensBefore, + TokensAfter: tokensBefore, + KeptCount: len(msgs), + NoOp: true, + }, nil + } + rebuilt := res.RebuildContext(msgs, now) + kept := len(rebuilt) - 1 + return &RebuildResult{ + Messages: rebuilt, + FromCheckpoint: false, + Watermark: res.FirstKeptIndex, + SummarizedCount: len(msgs) - kept, + KeptCount: kept, + TokensBefore: tokensBefore, + TokensAfter: compaction.EstimateContextTokens(rebuilt).Tokens, + }, nil +} + +// rebuildFromLoadedCheckpoint builds the rebuilt context from a loaded +// checkpoint: the pre-watermark prefix collapses to a single compaction message +// carrying cp.Summary, and the tail from the (clamped) watermark on is preserved +// verbatim. It reuses compaction.CompactionResult.RebuildContext so the summary +// message is shaped exactly like a compaction checkpoint. +func rebuildFromLoadedCheckpoint(msgs agentcore.MessageList, cp *Checkpoint, tokensBefore int, now int64) *RebuildResult { + w := cp.Watermark + if w < 0 { + w = 0 + } + if w > len(msgs) { + w = len(msgs) + } + res := &compaction.CompactionResult{ + Summary: cp.Summary, + FirstKeptIndex: w, + TokensBefore: tokensBefore, + } + rebuilt := res.RebuildContext(msgs, now) + kept := len(rebuilt) - 1 + return &RebuildResult{ + Messages: rebuilt, + FromCheckpoint: true, + Watermark: w, + SummarizedCount: w, + KeptCount: kept, + TokensBefore: tokensBefore, + TokensAfter: compaction.EstimateContextTokens(rebuilt).Tokens, + } +} diff --git a/pigo/internal/runtime/rebuild_test.go b/pigo/internal/runtime/rebuild_test.go new file mode 100644 index 0000000..ec4b863 --- /dev/null +++ b/pigo/internal/runtime/rebuild_test.go @@ -0,0 +1,194 @@ +package runtime + +// Tests for context rebuild (#482): with a checkpoint present, RebuildFromCheckpoint +// inserts the compression boundary at the watermark — collapsing the pre-watermark +// prefix into the checkpoint summary and preserving the recent tail verbatim; with +// no checkpoint it falls back to the lossy compaction path (compaction.Compact). +// Filesystem access uses t.TempDir(), matching checkpoint_test.go; the summary +// model is the shared summaryStream fake from compaction_test.go. + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/compaction" +) + +// textUser builds a user message carrying body, used to seed a rebuildable history. +func textUser(body string) agentcore.UserMessage { + return agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(body)}, + } +} + +func TestRebuildFromCheckpoint_InsertsBoundary(t *testing.T) { + root := t.TempDir() + const sessionID = "sess-rebuild" + + // A 6-message history; the checkpoint collapses the first 4 into a summary and + // keeps messages [4:] verbatim. + msgs := agentcore.MessageList{ + textUser("m0"), textUser("m1"), textUser("m2"), + textUser("m3"), textUser("keep-A"), textUser("keep-B"), + } + cp := Checkpoint{ + Watermark: 4, + Summary: "## Goal\ndistilled prefix", + CreatedAt: time.Now().UTC(), + CoveredMessages: 4, + } + if err := WriteCheckpoint(sessionID, root, cp); err != nil { + t.Fatalf("WriteCheckpoint: %v", err) + } + + // No summarization stream is needed: the checkpoint path is pure and local. + cfg := newRunCfg(nil) + res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil) + if err != nil { + t.Fatalf("RebuildFromCheckpoint: %v", err) + } + if !res.FromCheckpoint { + t.Fatalf("expected FromCheckpoint=true") + } + if res.NoOp { + t.Fatalf("expected a real rebuild, got NoOp") + } + if res.Watermark != 4 || res.SummarizedCount != 4 { + t.Fatalf("watermark/summarized: got %d/%d, want 4/4", res.Watermark, res.SummarizedCount) + } + // Rebuilt list = 1 compaction message + the retained tail (2 messages). + if len(res.Messages) != 3 { + t.Fatalf("rebuilt length: got %d, want 3: %+v", len(res.Messages), res.Messages) + } + if res.KeptCount != 2 { + t.Fatalf("kept: got %d, want 2", res.KeptCount) + } + // The prefix must collapse into a single compaction message carrying the summary. + head, ok := res.Messages[0].(agentcore.CompactionMessage) + if !ok { + t.Fatalf("message[0] should be a compaction checkpoint, got %T", res.Messages[0]) + } + if !strings.Contains(head.Summary, "distilled prefix") { + t.Fatalf("summary not carried through: %q", head.Summary) + } + // The recent tail is preserved verbatim, in order. + for i, want := range []string{"keep-A", "keep-B"} { + um, ok := res.Messages[i+1].(agentcore.UserMessage) + if !ok { + t.Fatalf("message[%d] should be preserved user message, got %T", i+1, res.Messages[i+1]) + } + if got := agentcore.ContentToText(um.Content); got != want { + t.Fatalf("tail[%d]: got %q, want %q", i, got, want) + } + } +} + +func TestRebuildFromCheckpoint_ClampsStaleWatermark(t *testing.T) { + root := t.TempDir() + const sessionID = "sess-stale" + + msgs := agentcore.MessageList{textUser("a"), textUser("b")} + // Watermark past the end (history was re-compacted since the checkpoint). + cp := Checkpoint{Watermark: 99, Summary: "old summary", CreatedAt: time.Now().UTC()} + if err := WriteCheckpoint(sessionID, root, cp); err != nil { + t.Fatalf("WriteCheckpoint: %v", err) + } + + cfg := newRunCfg(nil) + res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil) + if err != nil { + t.Fatalf("RebuildFromCheckpoint: %v", err) + } + // Clamped to len(msgs)=2: everything collapses, no verbatim tail, only the head. + if res.Watermark != 2 || res.KeptCount != 0 { + t.Fatalf("clamp: watermark=%d kept=%d, want 2/0", res.Watermark, res.KeptCount) + } + if len(res.Messages) != 1 { + t.Fatalf("rebuilt length: got %d, want 1 (summary only)", len(res.Messages)) + } +} + +func TestRebuildFromCheckpoint_WaitCallbackInvoked(t *testing.T) { + root := t.TempDir() + const sessionID = "sess-wait" + cp := Checkpoint{Watermark: 0, Summary: "s", CreatedAt: time.Now().UTC()} + if err := WriteCheckpoint(sessionID, root, cp); err != nil { + t.Fatalf("WriteCheckpoint: %v", err) + } + + waited := false + cfg := newRunCfg(nil) + _, err := RebuildFromCheckpoint(context.Background(), agentcore.MessageList{textUser("x")}, sessionID, root, &cfg, func() { waited = true }) + if err != nil { + t.Fatalf("RebuildFromCheckpoint: %v", err) + } + if !waited { + t.Fatalf("waitForCheckpoint callback was not invoked before reading") + } +} + +func TestRebuildFromCheckpoint_FallsBackToCompaction(t *testing.T) { + root := t.TempDir() // empty: no checkpoint on disk + const sessionID = "sess-nocp" + + // Seed a long history so FindCutPoint leaves a summarizable prefix. + msgs := bigUserMessages(12, 800) + + cfg := newRunCfg(scriptedStream(nil)) + cfg.SummaryStream = summaryStream("## Goal\nfallback compaction summary") + cfg.ContextWindow = 2000 + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + + res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil) + if err != nil { + t.Fatalf("RebuildFromCheckpoint: %v", err) + } + if res.FromCheckpoint { + t.Fatalf("expected fallback (FromCheckpoint=false) when no checkpoint exists") + } + if res.NoOp { + t.Fatalf("expected a real compaction fallback, got NoOp") + } + if res.SummarizedCount <= 0 { + t.Fatalf("expected some messages summarized, got %d", res.SummarizedCount) + } + if res.TokensAfter >= res.TokensBefore { + t.Fatalf("fallback should reduce tokens: before=%d after=%d", res.TokensBefore, res.TokensAfter) + } + // The rebuilt context begins with a compaction checkpoint holding the summary. + head, ok := res.Messages[0].(agentcore.CompactionMessage) + if !ok { + t.Fatalf("message[0] should be a compaction checkpoint, got %T", res.Messages[0]) + } + if !strings.Contains(head.Summary, "fallback compaction summary") { + t.Fatalf("fallback summary not carried through: %q", head.Summary) + } +} + +func TestRebuildFromCheckpoint_FallbackNoOpWhenNothingToCompact(t *testing.T) { + root := t.TempDir() // no checkpoint + const sessionID = "sess-empty" + + // A single short message: no valid cut point leaves a summarization range, so + // Compact returns (nil, nil) and the rebuild is a no-op. + msgs := agentcore.MessageList{textUser("only")} + cfg := newRunCfg(scriptedStream(nil)) + cfg.SummaryStream = summaryStream("unused") + cfg.ContextWindow = 2000 + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + + res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil) + if err != nil { + t.Fatalf("RebuildFromCheckpoint: %v", err) + } + if !res.NoOp { + t.Fatalf("expected NoOp when there is nothing to compact") + } + if len(res.Messages) != 1 || res.TokensAfter != res.TokensBefore { + t.Fatalf("no-op should return the original context unchanged: %+v", res) + } +} diff --git a/pigo/internal/runtime/reminder.go b/pigo/internal/runtime/reminder.go new file mode 100644 index 0000000..9c97fc9 --- /dev/null +++ b/pigo/internal/runtime/reminder.go @@ -0,0 +1,264 @@ +// This file implements the general system-reminder dynamic context injection +// mechanism (US-002, FR-1/FR-2), pigo's port of Claude Code's per-turn +// injection. +// +// A reminder is EPHEMERAL background context (the current todo list, a file +// that changed under the working directory, a budget warning) that should be +// visible to the model on the turn it matters, but must never pollute the +// durable conversation history. Two properties follow from that: +// +// - Not user instructions. Reminder bodies are wrapped in +// tags with a preamble stating they are background context from the harness, +// not a request from the user (the pi / Claude Code semantic convention). +// - Ephemeral. Reminders are injected only into the per-turn LLM request via +// the existing TransformContext seam, which shapes a COPY of the message +// list for the request and is never written back to AgentContext.Messages. +// Because they never enter the persisted message list they cannot be saved +// to the session file and cannot be folded into a compaction summary +// (compaction only ever sees AgentContext.Messages). +// +// The mechanism is a registry of ReminderProviders. Each provider is consulted +// every turn and may decline (ok == false) so a reminder only appears when its +// condition holds. RunConfig.Reminders wires the registry into the loop. +package runtime + +import ( + "context" + "strings" + "sync" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" +) + +// systemReminderPreamble marks the wrapped body as background context rather +// than a user instruction (FR-2). It leads every injected reminder so the model +// never mistakes harness state for a user request. +const systemReminderPreamble = "The following is background context provided automatically by the harness. " + + "It is NOT a message or instruction from the user; do not act on it as a request. " + + "Use it only to stay aware of the current state." + +// WrapSystemReminder wraps a reminder body in tags with the +// background-context preamble. The result is the text of a single injected +// message. +func WrapSystemReminder(body string) string { + return "\n" + systemReminderPreamble + "\n\n" + body + "\n" +} + +// ReminderProvider produces an ephemeral system-reminder for the upcoming turn. +// Reminder is consulted every turn with the current (post-TransformContext) +// message list; returning ok == false means "no reminder this turn", so a +// provider injects only when its condition holds. +type ReminderProvider interface { + // Name identifies the provider (for diagnostics/telemetry). It is not shown + // to the model. + Name() string + // Reminder returns the reminder body and true when a reminder should be + // injected this turn, or ("", false) to inject nothing. + Reminder(ctx context.Context, msgs agentcore.MessageList) (body string, ok bool) +} + +// ReminderFunc adapts a plain function to a ReminderProvider. +type ReminderFunc struct { + NameField string + Fn func(ctx context.Context, msgs agentcore.MessageList) (string, bool) +} + +// Name implements ReminderProvider. +func (f ReminderFunc) Name() string { return f.NameField } + +// Reminder implements ReminderProvider. +func (f ReminderFunc) Reminder(ctx context.Context, msgs agentcore.MessageList) (string, bool) { + if f.Fn == nil { + return "", false + } + return f.Fn(ctx, msgs) +} + +// ReminderRegistry holds the reminder providers consulted each turn. The zero +// value is usable (no providers → no injection); NewReminderRegistry is the +// convenience constructor. +type ReminderRegistry struct { + providers []ReminderProvider +} + +// NewReminderRegistry returns a registry pre-populated with providers. +func NewReminderRegistry(providers ...ReminderProvider) *ReminderRegistry { + r := &ReminderRegistry{} + for _, p := range providers { + r.Register(p) + } + return r +} + +// Register appends a provider. nil providers are ignored. +func (r *ReminderRegistry) Register(p ReminderProvider) { + if p == nil { + return + } + r.providers = append(r.providers, p) +} + +// Empty reports whether the registry has no providers (so callers can skip the +// injection wiring entirely). +func (r *ReminderRegistry) Empty() bool { return r == nil || len(r.providers) == 0 } + +// Messages consults every provider in registration order and returns the +// ephemeral reminder messages to inject this turn (one UserMessage per provider +// that fires). Reminders are modeled as user-role messages carrying +// -wrapped text — matching the pi / Claude Code convention +// where dynamic context enters through a user turn but is explicitly labeled as +// background context, not a user instruction. +func (r *ReminderRegistry) Messages(ctx context.Context, msgs agentcore.MessageList) []agentcore.AgentMessage { + if r.Empty() { + return nil + } + var out []agentcore.AgentMessage + for _, p := range r.providers { + body, ok := p.Reminder(ctx, msgs) + if !ok || body == "" { + continue + } + out = append(out, agentcore.UserMessage{ + RoleField: agentcore.RoleUser, + Content: agentcore.ContentList{agentcore.NewTextContent(WrapSystemReminder(body))}, + }) + } + return out +} + +// wrapTransform composes the registry into a TransformContext hook: it runs the +// caller's existing TransformContext (if any) first, then appends this turn's +// reminders to the shaped list. Because TransformContext output is used only to +// build the LLM request and is never written back to AgentContext.Messages, the +// appended reminders are ephemeral — they do not enter the persisted history and +// cannot be swept into a compaction summary. This is the single injection seam +// the loop wires in. +func (r *ReminderRegistry) wrapTransform( + inner func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList, +) func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList { + return func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList { + if inner != nil { + msgs = inner(ctx, msgs) + } + rem := r.Messages(ctx, msgs) + if len(rem) == 0 { + return msgs + } + out := make(agentcore.MessageList, 0, len(msgs)+len(rem)) + out = append(out, msgs...) + out = append(out, rem...) + return out + } +} + +// TodoReminderProvider is the built-in reference reminder provider (US-002): it +// surfaces the current todo list as background context whenever there is +// incomplete work, so the model is reminded of outstanding tasks each turn +// without the list having to be re-sent as a durable message. It reads the same +// TodoStore the todo tool writes, so the reminder always reflects the latest +// plan. When the list is empty or every item is completed it stays silent. +type TodoReminderProvider struct { + // Store is the session todo list. When nil the provider never fires. + Store *agenttool.TodoStore +} + +// Name implements ReminderProvider. +func (p *TodoReminderProvider) Name() string { return "todo" } + +// Reminder implements ReminderProvider. It fires only when the store holds at +// least one item that is not yet completed, keeping the condition deterministic +// and easy to test. +func (p *TodoReminderProvider) Reminder(ctx context.Context, _ agentcore.MessageList) (string, bool) { + if p.Store == nil { + return "", false + } + items := p.Store.Snapshot() + if len(items) == 0 { + return "", false + } + incomplete := false + for _, it := range items { + if it.Status != agenttool.TodoCompleted { + incomplete = true + break + } + } + if !incomplete { + return "", false + } + return "Your todo list has unfinished items. Keep it up to date with the todo tool.\n\n" + + agenttool.RenderTodoList(items), true +} + +// OneShotReminderProvider injects a fixed body on the NEXT turn only, then stays +// silent forever. Unlike the always-on providers (todo/goal) it does not depend +// on live state — it carries a snapshot of text captured at registration time. +// It exists for events that produce a single ephemeral injection, such as a +// UserPromptSubmit hook's additionalContext (US-007, FR-9): the hook's context +// must reach the model on the turn the prompt is sent, but must not persist into +// history or re-fire on later turns. sync.Once makes the single-fire transition +// safe even if the loop consults providers concurrently. +type OneShotReminderProvider struct { + name string + body string + once sync.Once +} + +// NewOneShotReminder builds a one-shot provider that will inject body exactly +// once. An empty body yields a provider that never fires. +func NewOneShotReminder(name, body string) *OneShotReminderProvider { + return &OneShotReminderProvider{name: name, body: body} +} + +// Name implements ReminderProvider. +func (p *OneShotReminderProvider) Name() string { + if p.name == "" { + return "one-shot" + } + return p.name +} + +// Reminder implements ReminderProvider. It returns its body and true on the very +// first consultation, then ("", false) on every subsequent turn. +func (p *OneShotReminderProvider) Reminder(ctx context.Context, _ agentcore.MessageList) (string, bool) { + if strings.TrimSpace(p.body) == "" { + return "", false + } + var body string + p.once.Do(func() { body = p.body }) + if body == "" { + return "", false + } + return body, true +} + +// GoalReminderProvider surfaces the active goal as background context each turn +// so the model keeps working toward it (mirrors pi-goal). It reads the same +// GoalState the /goal command drives, so the reminder always reflects the live +// objective. It fires only while the goal is active — a paused, blocked, or +// completed goal (and an idle state) injects nothing. +type GoalReminderProvider struct { + // State is the session goal state. When nil the provider never fires. + State *agenttool.GoalState +} + +// Name implements ReminderProvider. +func (p *GoalReminderProvider) Name() string { return "goal" } + +// Reminder implements ReminderProvider. It injects the objective plus a +// persistence instruction while the goal is active, and stays silent otherwise. +func (p *GoalReminderProvider) Reminder(ctx context.Context, _ agentcore.MessageList) (string, bool) { + if p.State == nil { + return "", false + } + snap := p.State.Snapshot() + if snap.Status != agenttool.GoalActive || strings.TrimSpace(snap.Objective) == "" { + return "", false + } + return "You are working autonomously toward this goal:\n\n" + snap.Objective + + "\n\nKeep making progress. When every requirement is verifiably met, call the " + + "goal_complete tool with a summary. If you hit a true impasse you cannot work " + + "around, call goal_blocked with concrete evidence. Do not stop or ask the user " + + "to continue — keep going until the goal is done or blocked.", true +} diff --git a/pigo/internal/runtime/reminder_test.go b/pigo/internal/runtime/reminder_test.go new file mode 100644 index 0000000..08c3ed0 --- /dev/null +++ b/pigo/internal/runtime/reminder_test.go @@ -0,0 +1,194 @@ +package runtime + +import ( + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/provider" +) + +// wrapStreamCapture wraps a StreamFn to record every text-block message the +// request carried into seen, so a test can assert what the model was sent. +func wrapStreamCapture(inner provider.StreamFn, seen *[]string) provider.StreamFn { + return func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + for _, m := range llm.Messages { + if um, ok := m.(agentcore.UserMessage); ok { + *seen = append(*seen, agentcore.ContentToText(um.Content)) + } + } + return inner(ctx, model, llm, cfg) + } +} + +// reminderTextsInRequest drives one turn and returns the +// message texts the provider stream actually received (the request-shaped list), +// so a test can assert what the model saw without touching persisted state. +func reminderTextsInRequest(t *testing.T, reg *ReminderRegistry, agentCtx *agentcore.AgentContext) []string { + t.Helper() + var seen []string + streamFn := scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, + }) + // Wrap the stream so we can inspect the LlmContext it is handed. + cfg := newRunCfg(nil) + cfg.Reminders = reg + cfg.LoopConfig.Stream = wrapStreamCapture(streamFn, &seen) + collectStream(t, agentLoop(context.Background(), agentCtx, cfg)) + return seen +} + +func TestWrapSystemReminderLabelsBackgroundContext(t *testing.T) { + out := WrapSystemReminder("body here") + if !strings.Contains(out, "") || !strings.Contains(out, "") { + t.Errorf("reminder must be wrapped in tags, got %q", out) + } + if !strings.Contains(out, "NOT a message or instruction from the user") { + t.Errorf("reminder must be labeled as background context, not a user instruction, got %q", out) + } + if !strings.Contains(out, "body here") { + t.Errorf("reminder must contain the body, got %q", out) + } +} + +func TestReminderInjectedWhenConditionHolds(t *testing.T) { + fired := ReminderFunc{NameField: "always", Fn: func(ctx context.Context, msgs agentcore.MessageList) (string, bool) { + return "budget is low", true + }} + reg := NewReminderRegistry(fired) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + seen := reminderTextsInRequest(t, reg, agentCtx) + found := false + for _, s := range seen { + if strings.Contains(s, "budget is low") && strings.Contains(s, "") { + found = true + } + } + if !found { + t.Fatalf("expected a system-reminder in the request, saw %v", seen) + } + // Ephemeral: the reminder must NOT be written back into the persisted history. + for _, m := range agentCtx.Messages { + if um, ok := m.(agentcore.UserMessage); ok { + if strings.Contains(agentcore.ContentToText(um.Content), "system-reminder") { + t.Errorf("reminder leaked into persisted message history: %+v", um) + } + } + } +} + +func TestReminderNotInjectedWhenConditionFails(t *testing.T) { + silent := ReminderFunc{NameField: "never", Fn: func(ctx context.Context, msgs agentcore.MessageList) (string, bool) { + return "", false + }} + reg := NewReminderRegistry(silent) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + seen := reminderTextsInRequest(t, reg, agentCtx) + for _, s := range seen { + if strings.Contains(s, "system-reminder") { + t.Errorf("no reminder should be injected when the provider declines, saw %q", s) + } + } +} + +func TestReminderPreservesInnerTransform(t *testing.T) { + var innerRan bool + reg := NewReminderRegistry(ReminderFunc{NameField: "always", Fn: func(ctx context.Context, msgs agentcore.MessageList) (string, bool) { + return "note", true + }}) + inner := func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList { + innerRan = true + return msgs + } + wrapped := reg.wrapTransform(inner) + out := wrapped(context.Background(), agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}) + if !innerRan { + t.Errorf("wrapTransform must call the inner TransformContext") + } + if len(out) != 2 { + t.Fatalf("expected original + 1 reminder message, got %d", len(out)) + } +} + +func TestTodoReminderProvider(t *testing.T) { + store := agenttool.NewTodoStore() + p := &TodoReminderProvider{Store: store} + + // Empty store: silent. + if _, ok := p.Reminder(context.Background(), nil); ok { + t.Errorf("empty todo store must not fire a reminder") + } + + // All completed: silent. + store.Set([]agenttool.TodoItem{{Content: "done", Status: agenttool.TodoCompleted}}) + if _, ok := p.Reminder(context.Background(), nil); ok { + t.Errorf("fully-completed todo list must not fire a reminder") + } + + // Incomplete work: fires with the rendered list. + store.Set([]agenttool.TodoItem{ + {Content: "write code", Status: agenttool.TodoInProgress}, + {Content: "review", Status: agenttool.TodoPending}, + }) + body, ok := p.Reminder(context.Background(), nil) + if !ok { + t.Fatalf("incomplete todo list must fire a reminder") + } + if !strings.Contains(body, "write code") || !strings.Contains(body, "review") { + t.Errorf("reminder body should render the todo items, got %q", body) + } + + // nil store: never fires. + nilP := &TodoReminderProvider{} + if _, ok := nilP.Reminder(context.Background(), nil); ok { + t.Errorf("nil todo store must not fire a reminder") + } +} + +func TestTodoReminderEndToEndInjection(t *testing.T) { + store := agenttool.NewTodoStore() + store.Set([]agenttool.TodoItem{{Content: "unfinished task", Status: agenttool.TodoInProgress}}) + reg := NewReminderRegistry(&TodoReminderProvider{Store: store}) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + seen := reminderTextsInRequest(t, reg, agentCtx) + found := false + for _, s := range seen { + if strings.Contains(s, "unfinished task") && strings.Contains(s, "") { + found = true + } + } + if !found { + t.Fatalf("todo reminder should be injected into the request, saw %v", seen) + } +} + +func TestGoalReminderOnlyWhenActive(t *testing.T) { + st := agenttool.NewGoalState() + p := &GoalReminderProvider{State: st} + + // Idle: no reminder. + if _, ok := p.Reminder(context.Background(), nil); ok { + t.Fatal("idle goal should not inject a reminder") + } + + // Active: injects the objective. + st.Start("g1", "build the feature", 0) + body, ok := p.Reminder(context.Background(), nil) + if !ok { + t.Fatal("active goal should inject a reminder") + } + if !strings.Contains(body, "build the feature") { + t.Errorf("reminder body missing objective: %q", body) + } + + // Completed: silent again. + st.MarkComplete("done") + if _, ok := p.Reminder(context.Background(), nil); ok { + t.Fatal("completed goal should not inject a reminder") + } +} diff --git a/pigo/internal/runtime/render.go b/pigo/internal/runtime/render.go new file mode 100644 index 0000000..174b80c --- /dev/null +++ b/pigo/internal/runtime/render.go @@ -0,0 +1,86 @@ +// This file implements DrainStream (architecture deepening ①): the single place +// that consumes a loop EventStream. The streaming-text delta accounting (track +// how many bytes of the current assistant message have been surfaced, emit only +// the new suffix) and the message_update-vs-turn_end dispatch were previously +// hand-rolled in three places — the REPL, the headless driver, and the +// sub-agent tool. Each consumer now supplies callbacks and shares one drain +// loop, so a bug in the delta arithmetic or the dispatch has exactly one home. +package runtime + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// StreamHandler is the set of callbacks DrainStream invokes as it consumes a +// run's events. Every field is optional (nil = ignore that signal). Callbacks +// run on the draining goroutine, in event order. +type StreamHandler struct { + // OnText receives each new suffix of the streaming assistant text: the bytes + // produced since the last OnText call for the current turn. The final suffix + // is flushed at turn end before OnTurnEnd, so a consumer that only implements + // OnText still sees the complete text. + OnText func(delta string) + // OnTurnEnd fires once per completed turn, after the turn's text is fully + // flushed, carrying the final assistant message and the tool results produced + // during the turn. Consumers render tool activity here. + OnTurnEnd func(msg agentcore.AssistantMessage, results []agentcore.ToolResultMessage) + // OnEvent, when set, receives every raw event before the typed callbacks — + // used by the stream-json protocol driver, which serialises the whole event. + OnEvent func(ev agentcore.AgentEvent) +} + +// DrainStream consumes stream to completion, invoking h's callbacks, and returns +// the final assistant message (or nil) plus the run's result error. It always +// drains every event even if a callback has side effects that fail, so the +// loop's producer goroutine never blocks on back-pressure (the no-leak +// contract). The final assistant message is taken from the run result when +// available, falling back to the last turn_end message observed on the stream. +func DrainStream(ctx context.Context, stream *LoopEventStream, h StreamHandler) (*agentcore.AssistantMessage, error) { + // printed tracks how many bytes of the current streaming assistant message + // have already been surfaced via OnText, so each update emits only the delta. + printed := 0 + var lastTurn *agentcore.AssistantMessage + + emitText := func(text string) { + if len(text) > printed { + if h.OnText != nil { + h.OnText(text[printed:]) + } + printed = len(text) + } + } + + for ev := range stream.Events() { + if h.OnEvent != nil { + h.OnEvent(ev) + } + switch e := ev.(type) { + case agentcore.MessageUpdateEvent: + if a, ok := e.Message.(agentcore.AssistantMessage); ok { + emitText(agentcore.ContentToText(a.Content)) + } + case agentcore.TurnEndEvent: + // Flush any tail the streaming updates did not cover (covers providers + // that only deliver the complete message at turn end), then reset for + // the next turn and hand the turn to the consumer. + emitText(agentcore.ContentToText(e.Message.Content)) + printed = 0 + m := e.Message + lastTurn = &m + if h.OnTurnEnd != nil { + h.OnTurnEnd(e.Message, e.ToolResults) + } + } + } + + msgs, resErr := stream.Result(ctx) + if resErr != nil { + return lastTurn, resErr + } + if final := agentcore.LastAssistantOf(msgs); final != nil { + return final, nil + } + return lastTurn, nil +} diff --git a/pigo/internal/runtime/render_test.go b/pigo/internal/runtime/render_test.go new file mode 100644 index 0000000..71c219a --- /dev/null +++ b/pigo/internal/runtime/render_test.go @@ -0,0 +1,164 @@ +package runtime + +// Tests for DrainStream (architecture deepening ①): the single event-stream +// consumer shared by the REPL, the headless driver, and the sub-agent tool. +// These drive it with a hand-built LoopEventStream so the delta accounting and +// the message_update-vs-turn_end dispatch are exercised directly, independent +// of any provider. + +import ( + "context" + "strings" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// emitStream builds a LoopEventStream, runs emit on a producer goroutine to +// push events and set the result, and returns the stream ready to drain. +func emitStream(emit func(s *LoopEventStream)) *LoopEventStream { + s := agentcore.NewEventStream[agentcore.AgentEvent, []agentcore.AgentMessage](0) + go func() { + emit(s) + s.Close() + }() + return s +} + +func assistantWith(text string) agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Content: agentcore.ContentList{agentcore.NewTextContent(text)}, + } +} + +// TestDrainStreamTextDeltas verifies OnText receives only the new suffix of the +// streaming assistant message on each update, never re-emitting printed bytes. +func TestDrainStreamTextDeltas(t *testing.T) { + ctx := context.Background() + stream := emitStream(func(s *LoopEventStream) { + _ = s.Emit(ctx, agentcore.MessageUpdateEvent{Message: assistantWith("Hel")}) + _ = s.Emit(ctx, agentcore.MessageUpdateEvent{Message: assistantWith("Hello")}) + _ = s.Emit(ctx, agentcore.MessageUpdateEvent{Message: assistantWith("Hello world")}) + _ = s.Emit(ctx, agentcore.TurnEndEvent{Message: assistantWith("Hello world")}) + s.SetResult([]agentcore.AgentMessage{assistantWith("Hello world")}) + }) + + var b strings.Builder + deltas := 0 + final, err := DrainStream(ctx, stream, StreamHandler{ + OnText: func(delta string) { b.WriteString(delta); deltas++ }, + }) + if err != nil { + t.Fatalf("DrainStream: %v", err) + } + if got := b.String(); got != "Hello world" { + t.Errorf("concatenated deltas = %q, want %q", got, "Hello world") + } + // 3 update deltas ("Hel","lo","world"); turn-end adds nothing new. + if deltas != 3 { + t.Errorf("OnText calls = %d, want 3 (turn-end must not re-emit)", deltas) + } + if final == nil || agentcore.ContentToText(final.Content) != "Hello world" { + t.Errorf("final message = %v, want %q", final, "Hello world") + } +} + +// TestDrainStreamTurnEndFlush verifies that when a provider delivers only the +// complete message at turn end (no streaming updates), OnText still receives the +// full text via the turn-end flush. +func TestDrainStreamTurnEndFlush(t *testing.T) { + ctx := context.Background() + stream := emitStream(func(s *LoopEventStream) { + _ = s.Emit(ctx, agentcore.TurnEndEvent{Message: assistantWith("complete only at end")}) + s.SetResult([]agentcore.AgentMessage{assistantWith("complete only at end")}) + }) + + var b strings.Builder + _, err := DrainStream(ctx, stream, StreamHandler{OnText: func(d string) { b.WriteString(d) }}) + if err != nil { + t.Fatalf("DrainStream: %v", err) + } + if got := b.String(); got != "complete only at end" { + t.Errorf("flushed text = %q, want full message", got) + } +} + +// TestDrainStreamResetsPerTurn verifies the delta accounting resets between +// turns, so a second turn's text is emitted from its own start rather than being +// masked by the first turn's printed offset. +func TestDrainStreamResetsPerTurn(t *testing.T) { + ctx := context.Background() + stream := emitStream(func(s *LoopEventStream) { + _ = s.Emit(ctx, agentcore.MessageUpdateEvent{Message: assistantWith("first")}) + _ = s.Emit(ctx, agentcore.TurnEndEvent{Message: assistantWith("first")}) + _ = s.Emit(ctx, agentcore.MessageUpdateEvent{Message: assistantWith("two")}) + _ = s.Emit(ctx, agentcore.TurnEndEvent{Message: assistantWith("two")}) + s.SetResult([]agentcore.AgentMessage{assistantWith("two")}) + }) + + var b strings.Builder + turns := 0 + _, err := DrainStream(ctx, stream, StreamHandler{ + OnText: func(d string) { b.WriteString(d) }, + OnTurnEnd: func(agentcore.AssistantMessage, []agentcore.ToolResultMessage) { turns++ }, + }) + if err != nil { + t.Fatalf("DrainStream: %v", err) + } + if got := b.String(); got != "firsttwo" { + t.Errorf("text across turns = %q, want %q", got, "firsttwo") + } + if turns != 2 { + t.Errorf("OnTurnEnd calls = %d, want 2", turns) + } +} + +// TestDrainStreamOnTurnEndResults verifies tool results are handed to OnTurnEnd. +func TestDrainStreamOnTurnEndResults(t *testing.T) { + ctx := context.Background() + res := agentcore.ToolResultMessage{Content: agentcore.ContentList{agentcore.NewTextContent("42")}} + stream := emitStream(func(s *LoopEventStream) { + _ = s.Emit(ctx, agentcore.TurnEndEvent{ + Message: assistantWith(""), + ToolResults: []agentcore.ToolResultMessage{res}, + }) + s.SetResult(nil) + }) + + var gotResults int + _, err := DrainStream(ctx, stream, StreamHandler{ + OnTurnEnd: func(_ agentcore.AssistantMessage, rs []agentcore.ToolResultMessage) { gotResults = len(rs) }, + }) + if err != nil { + t.Fatalf("DrainStream: %v", err) + } + if gotResults != 1 { + t.Errorf("OnTurnEnd tool results = %d, want 1", gotResults) + } +} + +// TestDrainStreamOnEvent verifies OnEvent sees every raw event in order (the +// stream-json driver's hook), and that a nil-callback handler still drains. +func TestDrainStreamOnEvent(t *testing.T) { + ctx := context.Background() + stream := emitStream(func(s *LoopEventStream) { + _ = s.Emit(ctx, agentcore.AgentStartEvent{}) + _ = s.Emit(ctx, agentcore.MessageUpdateEvent{Message: assistantWith("x")}) + _ = s.Emit(ctx, agentcore.TurnEndEvent{Message: assistantWith("x")}) + _ = s.Emit(ctx, agentcore.AgentEndEvent{}) + s.SetResult(nil) + }) + + var types []string + _, err := DrainStream(ctx, stream, StreamHandler{ + OnEvent: func(ev agentcore.AgentEvent) { types = append(types, ev.EventType()) }, + }) + if err != nil { + t.Fatalf("DrainStream: %v", err) + } + want := []string{"agent_start", "message_update", "turn_end", "agent_end"} + if strings.Join(types, ",") != strings.Join(want, ",") { + t.Errorf("OnEvent order = %v, want %v", types, want) + } +} diff --git a/pigo/internal/runtime/skills.go b/pigo/internal/runtime/skills.go new file mode 100644 index 0000000..8e4373c --- /dev/null +++ b/pigo/internal/runtime/skills.go @@ -0,0 +1,374 @@ +// This file implements declarative skills (US-028, #45): a skill is a markdown +// file with a YAML frontmatter block (mirrors SKILL.md) that names a reusable +// capability. Loading a skill parses its metadata (name, description, optional +// tool allow-list and model) and its markdown body (the skill's system prompt), +// then materializes it as a sub-agent tool — so invoking a skill runs its body +// as the system prompt of a child agent loop, reusing the SubAgentTool +// abstraction rather than introducing a second execution path. Each skill's +// description is surfaced in the parent's capability list so the model can +// choose to delegate to it. +package runtime + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" + "gopkg.in/yaml.v3" +) + +// SkillFrontmatter is the YAML metadata block at the head of a skill file. +type SkillFrontmatter struct { + // Name is the skill identifier; it becomes the spawnable tool name. When + // omitted it defaults to the file's base name (without extension). + Name string `yaml:"name"` + // Description tells the model what the skill does and when to use it. It is + // injected into the capability list, so it should be action-oriented. + Description string `yaml:"description"` + // AllowedTools optionally restricts the tools the skill's sub-agent may use, + // by tool name. Empty means "inherit the provided tool set as-is". Real + // Claude Code skills write this either as a YAML list or as a single + // scalar string (e.g. "Bash(foo:*), Read"), so it tolerates both forms. + AllowedTools stringList `yaml:"allowed-tools"` + // Model optionally pins the skill to a specific model; empty inherits. + Model string `yaml:"model"` + // DisableModelInvocation, when true, keeps the skill out of the system + // prompt's list so the model cannot auto-invoke it; it + // remains reachable only via its explicit "/name" slash command (mirrors pi's + // disable-model-invocation frontmatter key). Defaults to false. + DisableModelInvocation bool `yaml:"disable-model-invocation"` +} + +// Agent Skills spec limits (mirrors pi/agentskills.io): a skill name is a short +// slug and a description is a single sentence, both bounded so they stay cheap +// to inject into the system prompt. +const ( + maxSkillNameLength = 64 + maxSkillDescriptionLength = 1024 +) + +// skillNamePattern matches a valid skill name per the Agent Skills spec: +// lowercase ASCII letters, digits, and hyphens only. +var skillNamePattern = regexp.MustCompile(`^[a-z0-9-]+$`) + +// validateSkillName reports why name violates the Agent Skills spec, or nil +// when it is valid: lowercase a-z/0-9/hyphen only, at most maxSkillNameLength +// characters, and no leading, trailing, or consecutive hyphens. +func validateSkillName(name string) error { + if len(name) > maxSkillNameLength { + return fmt.Errorf("name exceeds %d characters (%d)", maxSkillNameLength, len(name)) + } + if !skillNamePattern.MatchString(name) { + return fmt.Errorf("name %q contains invalid characters (allowed: lowercase a-z, 0-9, hyphen)", name) + } + if strings.HasPrefix(name, "-") || strings.HasSuffix(name, "-") { + return fmt.Errorf("name %q must not start or end with a hyphen", name) + } + if strings.Contains(name, "--") { + return fmt.Errorf("name %q must not contain consecutive hyphens", name) + } + return nil +} + +// validateSkillDescription reports why description violates the spec, or nil +// when it is valid: non-empty and at most maxSkillDescriptionLength characters. +func validateSkillDescription(description string) error { + if strings.TrimSpace(description) == "" { + return errors.New("frontmatter missing required 'description'") + } + if len(description) > maxSkillDescriptionLength { + return fmt.Errorf("description exceeds %d characters (%d)", maxSkillDescriptionLength, len(description)) + } + return nil +} + +// stringList is a []string that unmarshals from either a YAML sequence +// (- a\n- b) or a single scalar. A scalar is split on commas so the common +// Claude Code form `allowed-tools: Bash(foo:*), Read` parses into two entries. +// This tolerance matters: a strict []string field rejects the scalar form and, +// because LoadSkillsDir aborts on the first parse error, one such skill would +// hide every other skill in the directory. +type stringList []string + +// UnmarshalYAML accepts a scalar or a sequence node. +func (l *stringList) UnmarshalYAML(node *yaml.Node) error { + switch node.Kind { + case yaml.ScalarNode: + var s string + if err := node.Decode(&s); err != nil { + return err + } + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + if t := strings.TrimSpace(p); t != "" { + out = append(out, t) + } + } + *l = out + return nil + case yaml.SequenceNode: + var ss []string + if err := node.Decode(&ss); err != nil { + return err + } + *l = ss + return nil + default: + // An empty/null node leaves the list nil (no restriction). + return nil + } +} + +// Skill is a parsed skill file: its metadata plus the markdown body that serves +// as the sub-agent's system prompt. +type Skill struct { + Frontmatter SkillFrontmatter + // Body is the markdown after the frontmatter block — the skill's instructions, + // used as the child agent's system prompt. + Body string + // Path is the source file, retained for diagnostics. + Path string +} + +// ParseSkill parses a skill's raw file content into a Skill. The file must open +// with a YAML frontmatter block delimited by lines containing only "---"; the +// remainder is the markdown body. A missing or malformed frontmatter block is +// an error, since name/description drive discovery. +func ParseSkill(path string, content []byte) (*Skill, error) { + fm, body, err := splitFrontmatter(content) + if err != nil { + return nil, fmt.Errorf("skill %s: %w", path, err) + } + var meta SkillFrontmatter + if err := yaml.Unmarshal(fm, &meta); err != nil { + return nil, fmt.Errorf("skill %s: parse frontmatter: %w", path, err) + } + if meta.Name == "" { + base := filepath.Base(path) + meta.Name = strings.TrimSuffix(base, filepath.Ext(base)) + } + if err := validateSkillName(meta.Name); err != nil { + return nil, fmt.Errorf("skill %s: %w", path, err) + } + if err := validateSkillDescription(meta.Description); err != nil { + return nil, fmt.Errorf("skill %s: %w", path, err) + } + return &Skill{Frontmatter: meta, Body: strings.TrimSpace(string(body)), Path: path}, nil +} + +// splitFrontmatter separates a leading "---"-delimited YAML block from the rest +// of the document. It returns the frontmatter bytes (without the fences) and +// the remaining body. It errors if the document does not open with a fence or +// the closing fence is missing. +func splitFrontmatter(content []byte) (frontmatter, body []byte, err error) { + text := string(content) + // Tolerate a UTF-8 BOM and leading blank lines before the opening fence. + text = strings.TrimPrefix(text, "\ufeff") + trimmed := strings.TrimLeft(text, "\r\n") + if !strings.HasPrefix(trimmed, "---") { + return nil, nil, fmt.Errorf("missing YAML frontmatter (file must start with '---')") + } + lines := strings.Split(trimmed, "\n") + // lines[0] is the opening fence. Find the closing fence. + var fmLines []string + closeIdx := -1 + for i := 1; i < len(lines); i++ { + if strings.TrimRight(lines[i], "\r") == "---" { + closeIdx = i + break + } + fmLines = append(fmLines, lines[i]) + } + if closeIdx == -1 { + return nil, nil, fmt.Errorf("unterminated YAML frontmatter (missing closing '---')") + } + bodyLines := lines[closeIdx+1:] + return []byte(strings.Join(fmLines, "\n")), []byte(strings.Join(bodyLines, "\n")), nil +} + +// LoadSkillsDir loads every "*.md" skill file in dir (non-recursively) plus any +// "/SKILL.md" nested layout (mirrors the SKILL.md convention). It returns the +// parsed skills sorted by name. A missing directory yields no skills and no +// error (skills are optional). +// +// A malformed skill file does NOT abort the load: the file is skipped and its +// error accumulated, so one bad skill cannot hide every other skill in the +// directory (a real ~/.agents/skills holds 100+ skills authored to varying +// conventions). The successfully parsed skills are always returned; the error, +// when non-nil, joins every skip reason for the caller to surface as a +// non-fatal warning. +func LoadSkillsDir(dir string) ([]*Skill, error) { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, fmt.Errorf("read skills dir %s: %w", dir, err) + } + var skills []*Skill + var errs []error + for _, e := range entries { + var path string + switch { + case e.IsDir(): + // Nested layout: //SKILL.md. + candidate := filepath.Join(dir, e.Name(), "SKILL.md") + if _, statErr := os.Stat(candidate); statErr != nil { + continue + } + path = candidate + case strings.EqualFold(filepath.Ext(e.Name()), ".md"): + path = filepath.Join(dir, e.Name()) + default: + continue + } + content, readErr := os.ReadFile(path) + if readErr != nil { + errs = append(errs, fmt.Errorf("read skill %s: %w", path, readErr)) + continue + } + skill, parseErr := ParseSkill(path, content) + if parseErr != nil { + errs = append(errs, parseErr) + continue + } + skills = append(skills, skill) + } + sort.Slice(skills, func(i, j int) bool { + return skills[i].Frontmatter.Name < skills[j].Frontmatter.Name + }) + return skills, errors.Join(errs...) +} + +// SubAgentSpec turns a skill into a sub-agent spec: the skill body becomes the +// child's system prompt, the description is surfaced to the model, and the tool +// set is the provided tools filtered by AllowedTools (when set). newRunConfig +// builds each child run's configuration; it receives the resolved tool set so +// the caller can wire a matching registry. +func (s *Skill) SubAgentSpec(tools []agentcore.AgentTool, newRunConfig func(tools []agentcore.AgentTool) RunConfig) SubAgentSpec { + resolved := filterToolsByName(tools, s.Frontmatter.AllowedTools) + return SubAgentSpec{ + Name: s.Frontmatter.Name, + Description: s.Frontmatter.Description, + SystemPrompt: s.Body, + Tools: resolved, + NewRunConfig: func() RunConfig { return newRunConfig(resolved) }, + } +} + +// SkillTool materializes a skill as an invocable sub-agent tool. +func (s *Skill) SkillTool(tools []agentcore.AgentTool, newRunConfig func(tools []agentcore.AgentTool) RunConfig) *SubAgentTool { + return NewSubAgentTool(s.SubAgentSpec(tools, newRunConfig)) +} + +// SlashCommand exposes the skill as a "/name" slash command (mirrors Claude Code's +// /skill-name invocation). Invoking it expands to the skill's instructions (its +// markdown body) as the prompt, with any arguments appended, so the skill runs +// in the current conversation. It is a prompt command (not an action): the +// expanded text is fed to the agent loop as the next user turn. +func (s *Skill) SlashCommand() SlashCommand { + body := s.Body + return SlashCommand{ + Name: s.Frontmatter.Name, + Description: s.Frontmatter.Description, + Source: SourceUser, + Expand: func(args string) string { + if strings.Contains(body, "$ARGUMENTS") { + return strings.ReplaceAll(body, "$ARGUMENTS", args) + } + if strings.TrimSpace(args) == "" { + return body + } + return body + "\n\n" + args + }, + } +} + +// FormatSkillsForPrompt renders the visible skills as an +// XML block for injection into the system prompt (mirrors pi's +// formatSkillsForPrompt). It implements progressive disclosure: only each +// skill's name, description, and location (the absolute SKILL.md path) are +// listed, so the model can read the file on demand rather than carrying every +// skill body in context. +// +// Skills with DisableModelInvocation == true are excluded (they remain +// reachable only via their explicit "/name" slash command). When no visible +// skill remains, it returns the empty string so callers can append +// unconditionally without altering a skill-free prompt. +func FormatSkillsForPrompt(skills []*Skill) string { + visible := make([]*Skill, 0, len(skills)) + for _, s := range skills { + if s != nil && !s.Frontmatter.DisableModelInvocation { + visible = append(visible, s) + } + } + if len(visible) == 0 { + return "" + } + lines := []string{ + "\n\nThe following skills provide specialized instructions for specific tasks.", + "Use the read tool to load a skill's file when the task matches its description.", + "When a skill file references a relative path, resolve it against the skill directory (parent of SKILL.md / dirname of the path) and use that absolute path in tool commands.", + "", + "", + } + for _, s := range visible { + lines = append(lines, + " ", + " "+escapeXML(s.Frontmatter.Name)+"", + " "+escapeXML(s.Frontmatter.Description)+"", + " "+escapeXML(skillLocation(s.Path))+"", + " ", + ) + } + lines = append(lines, "") + return strings.Join(lines, "\n") +} + +// skillLocation returns the absolute path to a skill file so the model can load +// it with the read tool regardless of the working directory. It falls back to +// the original path if resolution fails. +func skillLocation(path string) string { + if abs, err := filepath.Abs(path); err == nil { + return abs + } + return path +} + +// escapeXML escapes the five XML special characters so a skill's name or +// description cannot break the surrounding markup. +func escapeXML(s string) string { + r := strings.NewReplacer( + "&", "&", + "<", "<", + ">", ">", + `"`, """, + "'", "'", + ) + return r.Replace(s) +} + +// filterToolsByName keeps only tools whose Name is in allow. An empty allow +// list means "no restriction" and returns the input unchanged. +func filterToolsByName(tools []agentcore.AgentTool, allow []string) []agentcore.AgentTool { + if len(allow) == 0 { + return tools + } + set := make(map[string]bool, len(allow)) + for _, n := range allow { + set[n] = true + } + out := make([]agentcore.AgentTool, 0, len(tools)) + for _, t := range tools { + if set[t.Name()] { + out = append(out, t) + } + } + return out +} diff --git a/pigo/internal/runtime/slashcommand.go b/pigo/internal/runtime/slashcommand.go new file mode 100644 index 0000000..21a16d8 --- /dev/null +++ b/pigo/internal/runtime/slashcommand.go @@ -0,0 +1,538 @@ +// This file implements slash-commands (US-029, #45): typed "/name" shortcuts a +// user invokes in the REPL. There are two sources, resolved with a fixed +// priority: +// +// - Built-in commands are registered at compile time via RegisterBuiltin +// (from init() in the fork's own code). They are always available. +// - User commands are declarative markdown templates loaded from a directory +// (mirrors the .../commands/*.md convention): the file name is the command +// name and the body is a prompt template that may reference $ARGUMENTS. +// +// Conflict rule: same-name commands resolve by priority tier (built-in > +// project > global > package > settings > CLI); the higher tier wins and the +// loser is reported via Shadowed. Built-ins are load-bearing and always win. +// Within a tier, the last-added command overrides earlier ones (a re-load). +// +// There is deliberately no standalone plugin mechanism: a fork adds built-ins +// via init() registration, and external extensions go through MCP (deferred). +package runtime + +import ( + "fmt" + "os" + "path/filepath" + "sort" + "strings" + + "gopkg.in/yaml.v3" +) + +// SlashCommandSource identifies where a command came from, used for the +// conflict/priority rule and for display. +type SlashCommandSource int + +const ( + // SourceBuiltin is a compile-time registered command (highest priority). + SourceBuiltin SlashCommandSource = iota + // SourceUser is a declarative markdown command template loaded from disk + // (e.g. ~/.pigo/commands/*.md). + SourceUser + // SourceSkill is a skill loaded from ~/.agents/skills and surfaced as a + // /skill-name command. It behaves like SourceUser for the built-in-wins + // conflict rule; the finer tag is for display only (e.g. /status). + SourceSkill + // SourcePlugin is a command declared by a loaded plugin. It behaves like + // SourceUser for the built-in-wins conflict rule; the finer tag is for + // display only. + SourcePlugin +) + +func (s SlashCommandSource) String() string { + switch s { + case SourceBuiltin: + return "builtin" + case SourceSkill: + return "skill" + case SourcePlugin: + return "plugin" + default: + return "user" + } +} + +// Tier is the priority tier of a command, used to resolve same-name conflicts +// across sources (mirrors pi prompt-templates discovery priority). Higher tiers +// win; the loser is recorded in Shadowed. Within the same tier the last-added +// command wins (a re-load overrides). Skills and plugins are treated as +// Global-tier for priority - their finer Source label is for display only. +type Tier int + +const ( + // Tier values are ordered lowest-to-highest priority: in a same-name + // conflict the higher Tier value wins, so TierBuiltin always wins and + // TierCLI always loses. Declared ascending so the natural > comparison + // matches "higher priority wins". + TierCLI Tier = iota + // TierSettings is a prompt template referenced by the config.toml prompts array. + TierSettings + // TierPackage is a prompt template discovered from an installed package + // source (distinct from one copied into the global dir). + TierPackage + // TierGlobal is a global user prompt template (e.g. ~/.pigo/prompts or the + // legacy ~/.pigo/commands); also the tier used for skills and plugins. + TierGlobal + // TierProject is a project-local prompt template (e.g. .pigo/prompts). + TierProject + // TierBuiltin is a compile-time or instance built-in command (highest). + TierBuiltin +) + +func (t Tier) String() string { + switch t { + case TierBuiltin: + return "builtin" + case TierProject: + return "project" + case TierGlobal: + return "global" + case TierPackage: + return "package" + case TierSettings: + return "settings" + case TierCLI: + return "cli" + default: + return "unknown" + } +} + +// ShadowedEntry records a command that lost a same-name conflict to a higher- +// tier command, for diagnostics. It carries the loser's name, tier, and source +// label so /help and the startup warning can say which source was shadowed. +type ShadowedEntry struct { + Name string + Tier Tier + Source SlashCommandSource +} + +// String renders a shadowed entry as "name (tier)" for log lines. +func (e ShadowedEntry) String() string { return fmt.Sprintf("%s (%s)", e.Name, e.Tier) } + +// SlashCommand is a resolved command: its name (without the leading "/"), a +// short description for the command palette, and its source. A command is one +// of three kinds, distinguished by which callback is set: +// +// - A prompt command sets Expand: it turns the invocation arguments into the +// prompt text fed to the agent (the original slash-command behavior). +// - An action command sets Action instead: it performs a side effect (e.g. +// switching the runtime model) and returns a status line to show the user, +// rather than producing a prompt. No agent run is started. +// - A hybrid command sets Run: it performs a side effect AND may return prompt +// text to run — used by plugin commands, which RPC their plugin, surface the +// returned notifications, then inject the returned prompt as the next turn. +// +// Exactly one of Expand/Action/Run should be set. Precedence when more than one +// is set: Action wins over Run, which wins over Expand. This split is what lets +// a control command like "/model" change runtime state — the old design could +// only emit prompt text. +type SlashCommand struct { + Name string + Description string + // ArgumentHint is an optional frontmatter hint shown before the description + // in autocomplete (e.g. ""). Convention: for required args, + // [square] for optional. Empty when not set; display-only, not enforced. + ArgumentHint string + Source SlashCommandSource + // Tier is the priority tier used to resolve same-name conflicts across + // sources (built-in > project > global > package > settings > CLI). It is + // set by the AddX method matching the command's source; callers should not + // set it directly. + Tier Tier + // Expand maps the argument string (everything after "/name ") to the prompt + // text the command produces. For a built-in it may be arbitrary Go; for a + // user template it substitutes $ARGUMENTS into the markdown body. Nil for an + // action command. + Expand func(args string) string + // Action performs a side effect for the invocation and returns a status + // message to display (may be empty). Set instead of Expand for a control + // command like "/model". Because it is an arbitrary Go closure it can capture + // and mutate live runtime state, which Expand (a pure prompt producer) + // cannot. Nil for a prompt command. + Action func(args string) string + // Run is the hybrid of Action and Expand: it performs a side effect AND may + // produce prompt text to run as the next agent turn. It returns + // (message, prompt): message is shown to the user immediately (like an + // Action's status, e.g. plugin notifications), and prompt, when non-empty, is + // run as a normal turn (like Expand's output). This is what a plugin command + // needs — it RPCs its plugin (side effect), surfaces the returned + // notifications (message), then injects the returned prompt (prompt). Set + // instead of Expand/Action for such a command; nil otherwise. When Run is set + // it takes precedence over Expand (but Action still wins over Run). + Run func(args string) (message, prompt string) +} + +// SlashKind classifies how a resolved invocation should be handled by the +// caller: run its prompt through the agent, or treat it as a completed action. +type SlashKind int + +const ( + // SlashPrompt means the outcome carries prompt text to run (or, when not a + // command at all, the verbatim input). + SlashPrompt SlashKind = iota + // SlashAction means an action command already ran; the outcome carries only + // a status Message and no agent run should start. + SlashAction +) + +// SlashOutcome is the structured result of resolving one input line. Handled is +// false when the input was not a slash command (Prompt holds the verbatim input +// to run). When Handled is true, Kind says whether Prompt should be run +// (SlashPrompt) or an action already ran and Message should be shown without +// starting a run (SlashAction). +// +// A hybrid (Run) command resolves to Kind SlashPrompt with BOTH fields set: its +// side effect already ran, Message carries the text to show the user first +// (e.g. plugin notifications), and Prompt, when non-empty, is the turn to run +// after. The caller shows Message (if any) then runs Prompt (if non-empty). +type SlashOutcome struct { + Handled bool + Kind SlashKind + Prompt string + Message string +} + +// builtinCommands holds compile-time registered commands, keyed by name. It is +// populated by RegisterBuiltin from init() and read when building a registry. +// +// Concurrency contract: this global is written only by RegisterBuiltin, which +// must be called from init() (single-threaded, before main), and read only +// afterwards by NewSlashRegistry. It carries no lock because that init-only +// discipline means there is never a concurrent write; do not call +// RegisterBuiltin after startup. +var builtinCommands = map[string]SlashCommand{} + +// RegisterBuiltin registers a built-in slash command at compile time. It is +// intended to be called from init(); a duplicate name panics, since two +// built-ins claiming the same name is a programming error in the fork. +func RegisterBuiltin(cmd SlashCommand) { + if cmd.Name == "" { + panic("agent: RegisterBuiltin with empty name") + } + if _, exists := builtinCommands[cmd.Name]; exists { + panic(fmt.Sprintf("agent: duplicate built-in slash command %q", cmd.Name)) + } + cmd.Source = SourceBuiltin + cmd.Tier = TierBuiltin + builtinCommands[cmd.Name] = cmd +} + +// SlashRegistry resolves "/name" invocations against built-in and user +// commands, applying the built-in-wins priority rule. +type SlashRegistry struct { + commands map[string]SlashCommand + // shadowed records commands that lost a same-name conflict to a higher-tier + // command, with their tier and source for diagnostics. Same-tier overrides + // (last-write-wins) are not recorded. + shadowed []ShadowedEntry +} + +// NewSlashRegistry builds a registry seeded with all registered built-ins. +func NewSlashRegistry() *SlashRegistry { + r := &SlashRegistry{commands: make(map[string]SlashCommand, len(builtinCommands))} + for name, cmd := range builtinCommands { + r.commands[name] = cmd + } + return r +} + +// AddBuiltin installs a built-in command directly on this registry instance, +// bypassing the compile-time global. It exists for action commands whose +// closure must capture live, per-run state (e.g. a model controller created in +// main) — such state cannot be reached from an init()-time RegisterBuiltin. The +// command is marked SourceBuiltin so it wins over a same-named user command, +// exactly like a globally registered built-in. A duplicate name panics, since +// two built-ins claiming one name is a programming error. +func (r *SlashRegistry) AddBuiltin(cmd SlashCommand) { + if cmd.Name == "" { + panic("agent: AddBuiltin with empty name") + } + if existing, ok := r.commands[cmd.Name]; ok && existing.Source == SourceBuiltin { + panic(fmt.Sprintf("agent: duplicate built-in slash command %q", cmd.Name)) + } + cmd.Source = SourceBuiltin + cmd.Tier = TierBuiltin + r.add(cmd) +} + +// AddUser installs a user command (TierGlobal), e.g. a prompt template from +// ~/.pigo/prompts or the legacy ~/.pigo/commands. A same-named built-in or +// project-tier command wins; same-tier (global) adds override silently. +func (r *SlashRegistry) AddUser(cmd SlashCommand) { + cmd.Source = SourceUser + cmd.Tier = TierGlobal + r.add(cmd) +} + +// AddSkill installs a skill command (loaded from ~/.agents/skills) at TierGlobal. +// It follows the same tier rule as AddUser - a built-in or project-tier command +// wins - only the source tag differs, so /status can report skills separately. +func (r *SlashRegistry) AddSkill(cmd SlashCommand) { + cmd.Source = SourceSkill + cmd.Tier = TierGlobal + r.add(cmd) +} + +// AddPlugin installs a plugin-declared command at TierGlobal, mirroring AddUser +// with a SourcePlugin tag for display. +func (r *SlashRegistry) AddPlugin(cmd SlashCommand) { + cmd.Source = SourcePlugin + cmd.Tier = TierGlobal + r.add(cmd) +} + +// Shadowed returns the commands that lost a same-name conflict to a higher-tier +// command, with their tier and source for diagnostics. Same-tier overrides +// (last-write-wins) are not recorded here. +func (r *SlashRegistry) Shadowed() []ShadowedEntry { return r.shadowed } + +// AddProject installs a project-local prompt template (TierProject), which +// overrides a same-named global/package/settings/CLI template but loses to a +// built-in. +func (r *SlashRegistry) AddProject(cmd SlashCommand) { + cmd.Source = SourceUser + cmd.Tier = TierProject + r.add(cmd) +} + +// AddPackage installs a package-discovered prompt template (TierPackage). +func (r *SlashRegistry) AddPackage(cmd SlashCommand) { + cmd.Source = SourceUser + cmd.Tier = TierPackage + r.add(cmd) +} + +// AddSettings installs a prompt template referenced by config.toml (TierSettings). +func (r *SlashRegistry) AddSettings(cmd SlashCommand) { + cmd.Source = SourceUser + cmd.Tier = TierSettings + r.add(cmd) +} + +// AddCLI installs a prompt template referenced by --prompt-template (TierCLI, +// the lowest priority). +func (r *SlashRegistry) AddCLI(cmd SlashCommand) { + cmd.Source = SourceUser + cmd.Tier = TierCLI + r.add(cmd) +} + +// add installs cmd with tier-based conflict resolution. If a same-named command +// already exists, the higher tier wins and the loser is appended to shadowed; +// within the same tier the new command replaces the old (last-write-wins, no +// shadow entry). A built-in always wins because TierBuiltin is highest. +func (r *SlashRegistry) add(cmd SlashCommand) { + existing, ok := r.commands[cmd.Name] + if !ok { + r.commands[cmd.Name] = cmd + return + } + switch { + case existing.Tier > cmd.Tier: + // New command is lower tier: it loses and is shadowed. + r.shadowed = append(r.shadowed, ShadowedEntry{Name: cmd.Name, Tier: cmd.Tier, Source: cmd.Source}) + case existing.Tier < cmd.Tier: + // New command is higher tier: it wins; the old one is shadowed. + r.shadowed = append(r.shadowed, ShadowedEntry{Name: existing.Name, Tier: existing.Tier, Source: existing.Source}) + r.commands[cmd.Name] = cmd + default: + // Same tier: last-write-wins (a re-load), no shadow entry. + r.commands[cmd.Name] = cmd + } +} + +// Lookup returns the command bound to name (without the leading "/"). +func (r *SlashRegistry) Lookup(name string) (SlashCommand, bool) { + cmd, ok := r.commands[name] + return cmd, ok +} + +// List returns all commands sorted by name. +func (r *SlashRegistry) List() []SlashCommand { + out := make([]SlashCommand, 0, len(r.commands)) + for _, c := range r.commands { + out = append(out, c) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out +} + +// Resolve parses a raw input line and, if it is a slash-command invocation, +// expands it to the prompt text the agent should run. It returns (prompt, true) +// when input begins with "/" and names a known PROMPT command; (input, false) +// when the input is not a slash command (the caller runs it verbatim); and an +// error when input is a "/name" for an unknown command. +// +// This is the legacy string API, kept for callers that only handle prompt +// commands. It reports an action command as handled with an empty prompt (the +// action does NOT run here) — callers that want action commands to execute must +// use ResolveOutcome instead. +func (r *SlashRegistry) Resolve(input string) (prompt string, handled bool, err error) { + out, err := r.ResolveOutcome(input) + if err != nil { + return "", false, err + } + return out.Prompt, out.Handled, nil +} + +// ResolveOutcome parses a raw input line into a structured SlashOutcome. For a +// non-command it returns {Handled:false, Prompt:input}. For a known prompt +// command it returns {Handled:true, Kind:SlashPrompt, Prompt:}. For a +// known action command it RUNS the action and returns {Handled:true, +// Kind:SlashAction, Message:} — no prompt to run. For a known hybrid +// (Run) command it RUNS the side effect and returns {Handled:true, +// Kind:SlashPrompt, Message:, Prompt:} — the caller shows Message +// then runs Prompt when non-empty. An unknown "/name" yields an error. +func (r *SlashRegistry) ResolveOutcome(input string) (SlashOutcome, error) { + trimmed := strings.TrimLeft(input, " \t") + if !strings.HasPrefix(trimmed, "/") { + return SlashOutcome{Handled: false, Kind: SlashPrompt, Prompt: input}, nil + } + rest := trimmed[1:] + name := rest + args := "" + if i := strings.IndexAny(rest, " \t"); i >= 0 { + name = rest[:i] + args = strings.TrimSpace(rest[i+1:]) + } + cmd, ok := r.commands[name] + if !ok { + return SlashOutcome{}, fmt.Errorf("unknown command %q", "/"+name) + } + if cmd.Action != nil { + return SlashOutcome{Handled: true, Kind: SlashAction, Message: cmd.Action(args)}, nil + } + if cmd.Run != nil { + // A hybrid command runs its side effect now and may yield prompt text. + // The outcome is a prompt (SlashPrompt) that also carries a Message to + // surface first; the caller shows Message then runs Prompt if non-empty. + message, prompt := cmd.Run(args) + return SlashOutcome{Handled: true, Kind: SlashPrompt, Message: message, Prompt: prompt}, nil + } + return SlashOutcome{Handled: true, Kind: SlashPrompt, Prompt: cmd.Expand(args)}, nil +} + +// firstNonEmptyLine returns the first line of s whose trimmed form is non-empty, +// itself trimmed. It is the description fallback for templates whose frontmatter +// omits a description (mirrors pi: "If missing, the first non-empty line is used"). +func firstNonEmptyLine(s string) string { + for _, line := range strings.Split(s, "\n") { + if t := strings.TrimSpace(line); t != "" { + return t + } + } + return "" +} + +// LoadPromptFile loads a single prompt-template file. The command name is the +// filename without its extension (e.g. /x/review.md -> "review"). It is the +// single-file counterpart of LoadUserCommandsDir, used for settings/CLI paths +// that point at one file rather than a directory. +func LoadPromptFile(path string) (SlashCommand, error) { + content, err := os.ReadFile(path) + if err != nil { + return SlashCommand{}, fmt.Errorf("read prompt %s: %w", path, err) + } + base := filepath.Base(path) + name := strings.TrimSuffix(base, filepath.Ext(base)) + return ParseUserCommand(name, content) +} + +// LoadUserCommandsDir loads declarative markdown command templates from dir +// (non-recursively). Each "*.md" file defines a command named after the file +// (without extension). The file may carry an optional YAML frontmatter block +// with a "description" (mirrors skills); the remaining body is the prompt template, +// expanded via ExpandTemplate at invoke time. A missing directory yields no +// commands and no error. +func LoadUserCommandsDir(dir string) ([]SlashCommand, error) { + entries, err := os.ReadDir(dir) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, fmt.Errorf("read commands dir %s: %w", dir, err) + } + var cmds []SlashCommand + for _, e := range entries { + if e.IsDir() || !strings.EqualFold(filepath.Ext(e.Name()), ".md") { + continue + } + path := filepath.Join(dir, e.Name()) + content, readErr := os.ReadFile(path) + if readErr != nil { + return nil, fmt.Errorf("read command %s: %w", path, readErr) + } + name := strings.TrimSuffix(e.Name(), filepath.Ext(e.Name())) + cmd, parseErr := ParseUserCommand(name, content) + if parseErr != nil { + return nil, parseErr + } + cmds = append(cmds, cmd) + } + sort.Slice(cmds, func(i, j int) bool { return cmds[i].Name < cmds[j].Name }) + return cmds, nil +} + +// ParseUserCommand parses a declarative command template. An optional YAML +// frontmatter block supplies a description; the body is the prompt template, +// expanded at invoke time via ExpandTemplate (positional $N, $@/$ARGUMENTS, +// ${1:-default}, ${@:N}). If arg tokenization fails (e.g. an unterminated +// quote) the raw arg string is used as $ARGUMENTS so the invocation still works. +func ParseUserCommand(name string, content []byte) (SlashCommand, error) { + body := string(content) + description := "" + hint := "" + // Reuse the skills frontmatter splitter when a fence is present; otherwise + // treat the whole file as the template body. + if strings.HasPrefix(strings.TrimLeft(strings.TrimPrefix(body, "\ufeff"), "\r\n"), "---") { + fm, rest, splitErr := splitFrontmatter(content) + if splitErr != nil { + return SlashCommand{}, fmt.Errorf("command %s: %w", name, splitErr) + } + var meta struct { + Description string `yaml:"description"` + Name string `yaml:"name"` + ArgumentHint string `yaml:"argument-hint"` + } + if err := yaml.Unmarshal(fm, &meta); err != nil { + return SlashCommand{}, fmt.Errorf("command %s: parse frontmatter: %w", name, err) + } + description = meta.Description + hint = meta.ArgumentHint + if meta.Name != "" { + name = meta.Name + } + body = string(rest) + } + // When the frontmatter omits a description, fall back to the first non-empty + // line of the body (\u5bf9\u6807 pi: "If missing, the first non-empty line is used"). + if description == "" { + description = firstNonEmptyLine(body) + } + template := strings.TrimSpace(body) + return SlashCommand{ + Name: name, + Description: description, + ArgumentHint: hint, + Source: SourceUser, + Expand: func(args string) string { + tokens, err := SplitArgs(args) + if err != nil { + // Split failure (e.g. an unterminated quote): treat the raw arg + // string as a single $ARGUMENTS rather than feeding a malformed + // arg list to the engine, so a bad invocation stays usable. + tokens = []string{args} + } + return ExpandTemplate(template, tokens) + }, + }, nil +} diff --git a/pigo/internal/runtime/slashcommand_test.go b/pigo/internal/runtime/slashcommand_test.go new file mode 100644 index 0000000..c2af632 --- /dev/null +++ b/pigo/internal/runtime/slashcommand_test.go @@ -0,0 +1,182 @@ +package runtime + +// Tests for tiered priority resolution (US-009, #335): same-name commands +// across sources resolve by tier (built-in > project > global > package > +// settings > CLI), the loser is shadowed with its tier recorded, and same-tier +// adds override silently (last-write-wins, no shadow entry). + +import "testing" + +// TestSlashTierProjectOverridesGlobal: a project-tier template added after a +// global one wins; the global entry is shadowed with TierGlobal. +func TestSlashTierProjectOverridesGlobal(t *testing.T) { + r := NewSlashRegistry() + r.AddUser(SlashCommand{Name: "t", Expand: func(string) string { return "global" }}) + r.AddProject(SlashCommand{Name: "t", Expand: func(string) string { return "project" }}) + cmd, ok := r.Lookup("t") + if !ok { + t.Fatalf("command %q not found", "t") + } + if got := cmd.Expand(""); got != "project" { + t.Errorf("project must override global, got %q", got) + } + sh := r.Shadowed() + if len(sh) != 1 || sh[0].Name != "t" || sh[0].Tier != TierGlobal { + t.Errorf("global should be shadowed, got %v", sh) + } +} + +// TestSlashTierGlobalOverridesSettings: a global template added after a +// settings-tier one wins; the settings entry is shadowed with TierSettings. +func TestSlashTierGlobalOverridesSettings(t *testing.T) { + r := NewSlashRegistry() + r.AddSettings(SlashCommand{Name: "t", Expand: func(string) string { return "settings" }}) + r.AddUser(SlashCommand{Name: "t", Expand: func(string) string { return "global" }}) + cmd, ok := r.Lookup("t") + if !ok { + t.Fatalf("command %q not found", "t") + } + if got := cmd.Expand(""); got != "global" { + t.Errorf("global must override settings, got %q", got) + } + sh := r.Shadowed() + if len(sh) != 1 || sh[0].Name != "t" || sh[0].Tier != TierSettings { + t.Errorf("settings should be shadowed, got %v", sh) + } +} + +// TestSlashTierBuiltinOverridesProject: a built-in added after a project-tier +// template wins; the project entry is shadowed with TierProject. +func TestSlashTierBuiltinOverridesProject(t *testing.T) { + r := NewSlashRegistry() + r.AddProject(SlashCommand{Name: "t", Expand: func(string) string { return "project" }}) + r.AddBuiltin(SlashCommand{Name: "t", Action: func(string) string { return "builtin" }}) + cmd, ok := r.Lookup("t") + if !ok || cmd.Source != SourceBuiltin { + t.Fatalf("built-in must override project, got ok=%v source=%v", ok, cmd.Source) + } + sh := r.Shadowed() + if len(sh) != 1 || sh[0].Name != "t" || sh[0].Tier != TierProject { + t.Errorf("project should be shadowed, got %v", sh) + } +} + +// TestSlashTierSameTierLastWriteWins: two same-tier (global) adds resolve to the +// last one, with no shadow entry recorded. +func TestSlashTierSameTierLastWriteWins(t *testing.T) { + r := NewSlashRegistry() + r.AddUser(SlashCommand{Name: "t", Expand: func(string) string { return "first" }}) + r.AddUser(SlashCommand{Name: "t", Expand: func(string) string { return "second" }}) + cmd, ok := r.Lookup("t") + if !ok { + t.Fatalf("command %q not found", "t") + } + if got := cmd.Expand(""); got != "second" { + t.Errorf("same-tier last-write-wins, got %q", got) + } + if len(r.Shadowed()) != 0 { + t.Errorf("same-tier override must not shadow, got %v", r.Shadowed()) + } +} + +// TestSlashTierFullOrdering exercises the full tier ladder: adding lowest-first +// up to built-in, the built-in wins and every lower tier is shadowed. +func TestSlashTierFullOrdering(t *testing.T) { + r := NewSlashRegistry() + r.AddCLI(SlashCommand{Name: "t", Expand: func(string) string { return "cli" }}) + r.AddSettings(SlashCommand{Name: "t", Expand: func(string) string { return "settings" }}) + r.AddPackage(SlashCommand{Name: "t", Expand: func(string) string { return "package" }}) + r.AddUser(SlashCommand{Name: "t", Expand: func(string) string { return "global" }}) + r.AddProject(SlashCommand{Name: "t", Expand: func(string) string { return "project" }}) + r.AddBuiltin(SlashCommand{Name: "t", Action: func(string) string { return "builtin" }}) + cmd, ok := r.Lookup("t") + if !ok || cmd.Source != SourceBuiltin { + t.Fatalf("built-in must win the full ladder, got ok=%v source=%v", ok, cmd.Source) + } + // Five lower-tier commands (cli, settings, package, global, project) lost. + if len(r.Shadowed()) != 5 { + t.Errorf("expected 5 shadowed entries, got %d: %v", len(r.Shadowed()), r.Shadowed()) + } +} + +// Tests for ParseUserCommand wiring to the expansion engine (US-003, #333): +// Expand tokenizes args via SplitArgs and expands via ExpandTemplate, falling +// back to the raw arg string as $ARGUMENTS when tokenization fails. + +// TestParseUserCommandPositionalAndQuoted verifies multi-arg invocation, +// quoted-arg preservation, and the ${1:-default} form through ParseUserCommand. +func TestParseUserCommandPositionalAndQuoted(t *testing.T) { + // /review with no args: $ARGUMENTS expands to empty. + review, _ := ParseUserCommand("review", []byte("Review: $ARGUMENTS")) + if got := review.Expand(""); got != "Review: " { + t.Errorf("no args: got %q, want \"Review: \"", got) + } + // /component Button "click handler": quoted arg stays one token ($2). + comp, _ := ParseUserCommand("component", []byte("name=$1 feat=$2")) + if got := comp.Expand(`Button "click handler"`); got != "name=Button feat=click handler" { + t.Errorf("quoted args: got %q", got) + } + // ${1:-7} default: no arg -> 7, explicit -> the arg. + bul, _ := ParseUserCommand("summarize", []byte("in ${1:-7} bullets")) + if got := bul.Expand(""); got != "in 7 bullets" { + t.Errorf("default no arg: got %q", got) + } + if got := bul.Expand("5"); got != "in 5 bullets" { + t.Errorf("explicit arg: got %q", got) + } +} + +// TestParseUserCommandSplitFailureFallback verifies that an unterminated quote +// (SplitArgs error) falls back to treating the raw arg string as $ARGUMENTS. +func TestParseUserCommandSplitFailureFallback(t *testing.T) { + cmd, _ := ParseUserCommand("t", []byte("echo $ARGUMENTS")) + if got := cmd.Expand(`"unterminated`); got != `echo "unterminated` { + t.Errorf("split-failure fallback: got %q, want raw string as $ARGUMENTS", got) + } + // A no-placeholder template with split failure still appends the raw string. + bare, _ := ParseUserCommand("note", []byte("Take a note")) + if got := bare.Expand(`"unterminated`); got != "Take a note\n\n\"unterminated" { + t.Errorf("no-placeholder split failure: got %q", got) + } +} + +// TestParseUserCommandArgumentHintAndDescriptionFallback (US-004, #334): +// argument-hint is parsed from frontmatter, and description falls back to the +// first non-empty body line when absent. +func TestParseUserCommandArgumentHintAndDescriptionFallback(t *testing.T) { + // Both description and argument-hint. + cmd, err := ParseUserCommand("pr", []byte("---\ndescription: review PR\nargument-hint: \"\"\n---\nReview the PR")) + if err != nil { + t.Fatal(err) + } + if cmd.Description != "review PR" { + t.Errorf("description = %q, want \"review PR\"", cmd.Description) + } + if cmd.ArgumentHint != "" { + t.Errorf("argument-hint = %q, want \"\"", cmd.ArgumentHint) + } + // Only argument-hint: description falls back to first non-empty body line. + hintOnly, _ := ParseUserCommand("wr", []byte("---\nargument-hint: \"[instructions]\"\n---\nFinish the current task\nend-to-end")) + if hintOnly.Description != "Finish the current task" { + t.Errorf("description fallback = %q, want first body line", hintOnly.Description) + } + if hintOnly.ArgumentHint != "[instructions]" { + t.Errorf("argument-hint = %q, want \"[instructions]\"", hintOnly.ArgumentHint) + } + // Only description: argument-hint stays empty. + descOnly, _ := ParseUserCommand("cl", []byte("---\ndescription: audit changelog\n---\nAudit changelog entries")) + if descOnly.Description != "audit changelog" { + t.Errorf("description = %q", descOnly.Description) + } + if descOnly.ArgumentHint != "" { + t.Errorf("argument-hint should be empty, got %q", descOnly.ArgumentHint) + } + // No frontmatter at all: description falls back to first non-empty line. + neither, _ := ParseUserCommand("bare", []byte("First line is the desc\nSecond line is body")) + if neither.Description != "First line is the desc" { + t.Errorf("no-frontmatter fallback = %q, want first line", neither.Description) + } + if neither.ArgumentHint != "" { + t.Errorf("argument-hint should be empty, got %q", neither.ArgumentHint) + } +} diff --git a/pigo/internal/runtime/stream_response.go b/pigo/internal/runtime/stream_response.go new file mode 100644 index 0000000..3ab6c03 --- /dev/null +++ b/pigo/internal/runtime/stream_response.go @@ -0,0 +1,183 @@ +// This file implements streamAssistantResponse (US-003): it shapes the context +// into a provider request, resolves the API key dynamically, drives the +// provider stream, and back-fills the partial assistant message into the +// context while emitting message_start / message_update / message_end events. +package runtime + +import ( + "context" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/provider" +) + +// LoopConfig holds the pluggable behavior of the agent loop. Every hook is +// optional (nil = use the default). The pointer/func-field pattern mirrors pi's +// optional callbacks. +type LoopConfig struct { + // Model is the model id passed to StreamFn. + Model string + // APIKey is the static fallback key when GetAPIKey is nil or returns "". + APIKey string + // ThinkingLevel is the reasoning effort for requests. + ThinkingLevel agentcore.ThinkingLevel + // Stream produces the provider stream. Required (defaults are wired by + // callers/tests, e.g. a fake provider). + Stream provider.StreamFn + + // TransformContext optionally rewrites the message list before conversion + // (context trimming/injection). Contract: must not error; on failure return + // a safe fallback. Runs first. + TransformContext func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList + // ConvertToLlm optionally filters UI-only messages. Defaults to identity. + // Contract: must not error. + ConvertToLlm func(msgs agentcore.MessageList) agentcore.MessageList + // GetAPIKey optionally resolves a fresh key per request (handles short-lived + // token expiry). Falls back to APIKey when nil or empty. + GetAPIKey func(ctx context.Context, provider string) string + // Provider is the provider name passed to GetAPIKey. + Provider string + + // ContextWindow is the model's total context-token budget, used to decide + // automatic compaction. When <= 0 the window is unknown and auto-compaction + // is disabled (ShouldCompact returns false), so the loop behaves exactly as + // before for callers that do not plumb it through. + ContextWindow int + // Compaction holds the thresholds/retention knobs for auto-compaction. Its + // Enabled flag gates the feature independently of ContextWindow. + Compaction compaction.CompactionSettings + // SummaryStream produces the provider stream used to generate compaction + // summaries. Defaults to Stream when nil. + SummaryStream provider.StreamFn + // SummaryModel is the model used for summarization. When zero, a model is + // synthesized from Model/ContextWindow. + SummaryModel provider.Model + + // Extra is forwarded to StreamConfig.Extra. + Extra map[string]any +} + +// streamAssistantResponse runs one assistant turn: it builds the request from +// agentCtx, streams the provider response, back-fills the partial into +// agentCtx.Messages, and returns the final assistant message. The sequence +// (transformContext → convertToLlm → resolve key → stream → drain) is kept +// identical to pi. It never returns an error for a request failure — such +// failures arrive as a terminal assistant message with stopReason error/aborted. +func streamAssistantResponse(ctx context.Context, agentCtx *agentcore.AgentContext, cfg LoopConfig, emit agentcore.EmitFunc) (agentcore.AssistantMessage, error) { + // 1. transformContext (optional, must not error). + msgs := agentCtx.Messages + if cfg.TransformContext != nil { + msgs = cfg.TransformContext(ctx, msgs) + } + // 2. convertToLlm (filter UI-only; default identity). + if cfg.ConvertToLlm != nil { + msgs = cfg.ConvertToLlm(msgs) + } + // 3. shape the LLM context. + llm := provider.LlmContext{ + SystemPrompt: agentCtx.SystemPrompt, + Messages: msgs, + Tools: agentCtx.Tools, + } + // 4. resolve API key dynamically, fall back to static. + key := cfg.APIKey + if cfg.GetAPIKey != nil { + if dyn := cfg.GetAPIKey(ctx, cfg.Provider); dyn != "" { + key = dyn + } + } + // 5. build the provider stream. + stream, err := cfg.Stream(ctx, cfg.Model, llm, provider.StreamConfig{ + APIKey: key, + ThinkingLevel: cfg.ThinkingLevel, + Extra: cfg.Extra, + }) + if err != nil { + // Early "cannot build stream" failure: synthesize a terminal message so + // the loop has a uniform assistant message to record. + return newErrorAssistantMessage(cfg, err), nil + } + + // 6. drain the stream, back-filling the partial into the context. + addedPartial := false + backfill := func(partial agentcore.AssistantMessage) { + if !addedPartial { + agentCtx.Messages = append(agentCtx.Messages, partial) + addedPartial = true + } else { + agentCtx.Messages[len(agentCtx.Messages)-1] = partial + } + } + + for ev := range stream.Events() { + switch e := ev.(type) { + case provider.StreamStartEvent: + backfill(e.Partial) + if err := emit(ctx, agentcore.MessageStartEvent{Message: e.Partial}); err != nil { + return agentcore.AssistantMessage{}, err + } + case provider.StreamTextEvent: + backfill(e.Partial) + if err := emit(ctx, agentcore.MessageUpdateEvent{Message: e.Partial, AssistantMessageEvent: e}); err != nil { + return agentcore.AssistantMessage{}, err + } + case provider.StreamThinkingEvent: + backfill(e.Partial) + if err := emit(ctx, agentcore.MessageUpdateEvent{Message: e.Partial, AssistantMessageEvent: e}); err != nil { + return agentcore.AssistantMessage{}, err + } + case provider.StreamToolCallEvent: + backfill(e.Partial) + if err := emit(ctx, agentcore.MessageUpdateEvent{Message: e.Partial, AssistantMessageEvent: e}); err != nil { + return agentcore.AssistantMessage{}, err + } + case provider.StreamDoneEvent: + finalizeMessage(agentCtx, e.Message, &addedPartial) + if err := emit(ctx, agentcore.MessageEndEvent{Message: e.Message}); err != nil { + return agentcore.AssistantMessage{}, err + } + return e.Message, nil + case provider.StreamErrorEvent: + finalizeMessage(agentCtx, e.Message, &addedPartial) + if err := emit(ctx, agentcore.MessageEndEvent{Message: e.Message}); err != nil { + return agentcore.AssistantMessage{}, err + } + return e.Message, nil + } + } + + // 7. stream ended without done/error: fall back to the stream result. + final, resErr := stream.Result(ctx) + if resErr != nil { + return newErrorAssistantMessage(cfg, resErr), nil + } + finalizeMessage(agentCtx, final, &addedPartial) + if err := emit(ctx, agentcore.MessageEndEvent{Message: final}); err != nil { + return agentcore.AssistantMessage{}, err + } + return final, nil +} + +// finalizeMessage replaces the placeholder partial with the final message, or +// appends it if the provider sent done/error without a prior start. +func finalizeMessage(agentCtx *agentcore.AgentContext, final agentcore.AssistantMessage, addedPartial *bool) { + if *addedPartial { + agentCtx.Messages[len(agentCtx.Messages)-1] = final + } else { + agentCtx.Messages = append(agentCtx.Messages, final) + *addedPartial = true + } +} + +// newErrorAssistantMessage builds a terminal assistant message for an early +// failure that never produced a provider stream. +func newErrorAssistantMessage(cfg LoopConfig, err error) agentcore.AssistantMessage { + return agentcore.AssistantMessage{ + RoleField: agentcore.RoleAssistant, + Model: cfg.Model, + Provider: cfg.Provider, + StopReason: agentcore.StopReasonError, + ErrorMessage: err.Error(), + } +} diff --git a/pigo/internal/runtime/stream_response_test.go b/pigo/internal/runtime/stream_response_test.go new file mode 100644 index 0000000..4e6abc7 --- /dev/null +++ b/pigo/internal/runtime/stream_response_test.go @@ -0,0 +1,176 @@ +package runtime + +import ( + "context" + "testing" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// fakeStream builds a StreamFn that replays a fixed sequence of events, pushing +// each onto an AssistantMessageEventStream from a producer goroutine. +func fakeStream(events []provider.AssistantMessageEvent) provider.StreamFn { + return func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + s := provider.NewAssistantMessageEventStream(0) + go func() { + for _, ev := range events { + if err := s.Emit(ctx, ev); err != nil { + s.SetError(err) + s.Close() + return + } + } + s.Close() + }() + return s, nil + } +} + +// drives streamAssistantResponse with a synchronous emit that records events. +func runStream(t *testing.T, agentCtx *agentcore.AgentContext, cfg LoopConfig) (agentcore.AssistantMessage, []agentcore.AgentEvent) { + t.Helper() + var got []agentcore.AgentEvent + emit := func(ctx context.Context, ev agentcore.AgentEvent) error { + got = append(got, ev) + return nil + } + msg, err := streamAssistantResponse(context.Background(), agentCtx, cfg, emit) + if err != nil { + t.Fatalf("streamAssistantResponse: %v", err) + } + return msg, got +} + +func TestStreamResponseBackfillAndEvents(t *testing.T) { + partial0 := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + partial1 := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("hel")}} + final := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("hello")}, StopReason: agentcore.StopReasonEndTurn} + + cfg := LoopConfig{ + Model: "fake", + Stream: fakeStream([]provider.AssistantMessageEvent{ + provider.StreamStartEvent{Partial: partial0}, + provider.StreamTextEvent{Partial: partial1}, + provider.StreamDoneEvent{Message: final}, + }), + } + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + msg, events := runStream(t, agentCtx, cfg) + + if msg.StopReason != agentcore.StopReasonEndTurn { + t.Errorf("final stopReason = %q, want end_turn", msg.StopReason) + } + // Context should hold the user message + the final assistant message (the + // placeholder was replaced, not appended twice). + if len(agentCtx.Messages) != 2 { + t.Fatalf("context messages = %d, want 2: %+v", len(agentCtx.Messages), agentCtx.Messages) + } + last, ok := agentCtx.Messages[1].(agentcore.AssistantMessage) + if !ok || len(last.Content) != 1 { + t.Fatalf("last message not final assistant: %+v", agentCtx.Messages[1]) + } + // Event order: message_start, message_update, message_end. + wantKinds := []string{agentcore.EventMessageStart, agentcore.EventMessageUpdate, agentcore.EventMessageEnd} + if len(events) != len(wantKinds) { + t.Fatalf("event count = %d, want %d: %+v", len(events), len(wantKinds), events) + } + for i, w := range wantKinds { + if events[i].EventType() != w { + t.Errorf("event[%d] = %q, want %q", i, events[i].EventType(), w) + } + } +} + +func TestStreamResponseErrorEvent(t *testing.T) { + errMsg := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonError, ErrorMessage: "boom"} + cfg := LoopConfig{ + Model: "fake", + Stream: fakeStream([]provider.AssistantMessageEvent{provider.StreamErrorEvent{Message: errMsg}}), + } + agentCtx := &agentcore.AgentContext{} + msg, events := runStream(t, agentCtx, cfg) + if msg.StopReason != agentcore.StopReasonError || msg.ErrorMessage != "boom" { + t.Errorf("want error terminal message, got %+v", msg) + } + // No start event was sent; error should still append the terminal message. + if len(agentCtx.Messages) != 1 { + t.Fatalf("context messages = %d, want 1", len(agentCtx.Messages)) + } + if events[len(events)-1].EventType() != agentcore.EventMessageEnd { + t.Errorf("last event = %q, want message_end", events[len(events)-1].EventType()) + } +} + +func TestStreamResponseDynamicAPIKey(t *testing.T) { + var seenKey string + streamFn := func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + seenKey = cfg.APIKey + s := provider.NewAssistantMessageEventStream(0) + go func() { + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}}) + s.Close() + }() + return s, nil + } + cfg := LoopConfig{ + Model: "fake", + APIKey: "static-key", + Provider: "test", + Stream: streamFn, + GetAPIKey: func(ctx context.Context, provider string) string { return "dynamic-key" }, + } + runStream(t, &agentcore.AgentContext{}, cfg) + if seenKey != "dynamic-key" { + t.Errorf("dynamic key not used: got %q", seenKey) + } + + // Empty dynamic key falls back to static. + cfg.GetAPIKey = func(ctx context.Context, provider string) string { return "" } + runStream(t, &agentcore.AgentContext{}, cfg) + if seenKey != "static-key" { + t.Errorf("fallback to static key failed: got %q", seenKey) + } +} + +func TestStreamResponseTransformAndConvertOrder(t *testing.T) { + var order []string + cfg := LoopConfig{ + Model: "fake", + TransformContext: func(ctx context.Context, msgs agentcore.MessageList) agentcore.MessageList { + order = append(order, "transform") + return msgs + }, + ConvertToLlm: func(msgs agentcore.MessageList) agentcore.MessageList { + order = append(order, "convert") + return msgs + }, + Stream: func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + order = append(order, "stream") + s := provider.NewAssistantMessageEventStream(0) + go func() { + _ = s.Emit(ctx, provider.StreamDoneEvent{Message: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}}) + s.Close() + }() + return s, nil + }, + } + runStream(t, &agentcore.AgentContext{}, cfg) + if len(order) != 3 || order[0] != "transform" || order[1] != "convert" || order[2] != "stream" { + t.Errorf("call order wrong: %v", order) + } +} + +func TestStreamResponseEarlyBuildFailure(t *testing.T) { + cfg := LoopConfig{ + Model: "fake", + Stream: func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + return nil, context.DeadlineExceeded + }, + } + msg, _ := runStream(t, &agentcore.AgentContext{}, cfg) + if msg.StopReason != agentcore.StopReasonError { + t.Errorf("early build failure should yield error message, got %+v", msg) + } +} diff --git a/pigo/internal/runtime/subagent.go b/pigo/internal/runtime/subagent.go new file mode 100644 index 0000000..e67a7d4 --- /dev/null +++ b/pigo/internal/runtime/subagent.go @@ -0,0 +1,472 @@ +// This file implements sub-agent orchestration (US-027, #45) with an optional +// process-isolation mode (US-019, #135). +// +// A sub-agent is a full agent loop with its own AgentContext (independent system +// prompt, message history and tool set), launched by the parent through a normal +// tool call. The child runs to completion and its final assistant text is fed +// back to the parent as the tool result - so from the parent loop's perspective +// a sub-agent is just another tool. +// +// Two isolation modes are supported, selected by SubAgentSpec.Isolation: +// +// - Goroutine (default): the child loop runs in-process in a goroutine sharing +// the parent process, matching the original "single-process goroutine" decision. +// - Process: the parent spawns a fresh pigo subprocess (pigo --subagent-rpc) +// and delegates the run over stdio JSON-RPC (reusing internal/jsonrpc). The +// child runs in a separate process, so a crash or resource leak in the child +// cannot affect the parent loop; a crash is surfaced as a tool error. The +// subprocess resolves its own provider from the model/provider passed in the +// request and inherits the parent environment for credentials. +// +// Because each Execute call spins up an independent run, multiple sub-agents can +// run concurrently (the batch executor already runs parallel tool calls in +// separate goroutines/processes). +package runtime + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/jsonrpc" +) + +// SubAgentIsolation selects how a sub-agent runs relative to its parent. +type SubAgentIsolation int + +const ( + // SubAgentIsolationGoroutine runs the child agent loop in-process in a + // goroutine. This is the default and the original behavior; it changes + // nothing about how sub-agents previously ran. + SubAgentIsolationGoroutine SubAgentIsolation = iota + // SubAgentIsolationProcess runs the child in a fresh pigo subprocess, + // delegating the run over stdio JSON-RPC. A subprocess crash is surfaced to + // the parent as a tool error and never affects the parent loop. + SubAgentIsolationProcess +) + +// SubAgentProcessConfig configures process-isolated sub-agent execution. It +// carries the serializable provider config the subprocess needs to reconstruct +// the run: the in-process Stream/GetAPIKey functions a goroutine-mode +// NewRunConfig returns cannot cross a process boundary, so the parent forwards +// the model (and optional base URL/protocol) and the subprocess resolves the +// provider itself, inheriting the parent environment for API keys. +type SubAgentProcessConfig struct { + // Command is the executable to spawn. When empty, os.Executable() (the pigo + // binary itself) is used so a pigo process spawns another pigo. + Command string + // Args are appended to the command after the subagent-rpc flag. Rarely + // needed; reserved for test doubles or non-standard layouts. + Args []string + // Model is the model id the subprocess runs against. Required. A preset id + // (e.g. "openrouter/free", "anthropic/claude-...") or ollama/nvidia-prefixed + // id resolves its own provider; a custom gateway needs BaseURL/Protocol. + Model string + // BaseURL and Protocol override the provider endpoint and wire protocol for + // custom gateways (Protocol "anthropic"/"openai" forces that wire format). + // Empty falls back to the same resolution the CLI uses. + BaseURL string + Protocol string + // ToolNames restricts the subprocess's builtin tool set to the named tools + // (e.g. a read-only researcher). Empty keeps all builtins. Non-builtin names + // are ignored: custom/plugin tools cannot cross a process boundary, so a + // process-isolated child runs with builtins only. + ToolNames []string + // Env is the child's environment (os/exec form). When nil the child inherits + // the parent environment, which is how it picks up provider API keys. + Env []string + // Dir is the child's working directory; empty means the parent's. + Dir string + // Stderr optionally receives the child's stderr. When nil it is discarded. + Stderr io.Writer +} + +// SubAgentRunParams is the JSON-RPC request payload for a process-isolated +// sub-agent run (the "subagent/run" method). It is the wire contract between +// the parent (SubAgentTool in process mode) and the pigo subprocess +// (cmd/pigo --subagent-rpc). +type SubAgentRunParams struct { + Prompt string `json:"prompt"` + SystemPrompt string `json:"systemPrompt,omitempty"` + Model string `json:"model"` + BaseURL string `json:"baseUrl,omitempty"` + Protocol string `json:"protocol,omitempty"` + Tools []string `json:"tools,omitempty"` +} + +// SubAgentRunResult is the JSON-RPC response payload carrying the child's final +// assistant text. +type SubAgentRunResult struct { + Text string `json:"text"` +} + +// SubAgentRPCMethod is the JSON-RPC method name the parent calls on the +// subprocess: "subagent/run". +const SubAgentRPCMethod = "subagent/run" + +// SubAgentRPCFlag is the command-line flag the parent launches the pigo +// subprocess with so it enters the sub-agent RPC server mode: "--subagent-rpc". +const SubAgentRPCFlag = "--subagent-rpc" + +// SubAgentSpec declares a spawnable sub-agent: its identity (surfaced to the +// model as a tool), the system prompt and tools its child context runs with, +// and a factory for the child's run configuration (provider stream, batch +// registry, hooks). The factory is called once per spawn so each child gets an +// independent RunConfig; NewRunConfig must wire a ToolRegistry consistent with +// Tools. It is used by goroutine mode; process mode uses Process instead (the +// subprocess builds its own RunConfig from the serializable provider config). +type SubAgentSpec struct { + // Name is the tool name the parent invokes to spawn this sub-agent. + Name string + // Description is injected into the parent's tool list / capability list so + // the model knows when to delegate. + Description string + // SystemPrompt seeds the child context's system prompt. When empty the child + // runs with no system prompt. + SystemPrompt string + // Tools is the child's independent tool set. It may differ from the parent's + // (e.g. a read-only researcher sub-agent) and may be empty. In goroutine + // mode these exact tools run in-process; in process mode only the tools' + // NAMES are forwarded (the subprocess rebuilds builtins by name). + Tools []agentcore.AgentTool + // NewRunConfig builds the loop configuration for one child run in goroutine + // mode. It is called per spawn; the returned config's Batch registry should + // contain Tools. Ignored in process mode (the subprocess resolves its own). + NewRunConfig func() RunConfig + // Isolation selects goroutine (default) vs process execution. Zero value is + // goroutine, preserving the original behavior. + Isolation SubAgentIsolation + // Process configures process-isolated execution. Required when Isolation is + // SubAgentIsolationProcess; ignored otherwise. + Process SubAgentProcessConfig + // Schema, when non-empty, overrides the default single-prompt argument schema + // advertised to the model. The generic task tool uses this to also accept an + // optional description; a nil/empty Schema keeps the original prompt-only + // schema so existing specs are unaffected. + Schema json.RawMessage + // Sem, when non-nil, is a shared buffered channel used as a concurrency + // semaphore for goroutine-mode runs: executeGoroutine acquires a slot before + // spawning the child and releases it when the child settles. A full channel + // blocks (queues) the acquire rather than erroring. nil disables limiting, so + // existing sub-agent specs run unbounded exactly as before. + Sem chan struct{} +} + +// subAgentArgs is the JSON argument shape for a sub-agent tool call: a +// free-form prompt describing the delegated task, plus an optional short +// description used for status display (accepted by the generic task tool; +// ignored by prompt-only specs). +type subAgentArgs struct { + Prompt string `json:"prompt"` + Description string `json:"description,omitempty"` +} + +// subAgentSchema is the JSON Schema validating a sub-agent invocation. +var subAgentSchema = json.RawMessage(`{ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "The task for the sub-agent to perform, described in full since the sub-agent runs with a fresh context." + } + }, + "required": ["prompt"], + "additionalProperties": false +}`) + +// SubAgentTool adapts a SubAgentSpec into an AgentTool. Executing it spawns a +// child agent run (in a goroutine or a subprocess, per Isolation) and returns +// the child's final text. +type SubAgentTool struct { + spec SubAgentSpec + // processCall, when non-nil, overrides the default subprocess transport for + // process-isolated mode. Tests inject a fake to exercise the process-mode + // logic (params shaping, crash-as-error, result forwarding) without building + // a real binary; production leaves it nil so Execute uses defaultProcessCall. + processCall func(ctx context.Context, cfg SubAgentProcessConfig, params SubAgentRunParams) (string, error) +} + +// NewSubAgentTool builds a sub-agent tool from a spec. In goroutine mode +// NewRunConfig is required (it supplies the provider stream that drives the +// child); in process mode Process.Model is required instead. +func NewSubAgentTool(spec SubAgentSpec) *SubAgentTool { + return &SubAgentTool{spec: spec} +} + +func (t *SubAgentTool) Name() string { return t.spec.Name } + +func (t *SubAgentTool) Description() string { return t.spec.Description } + +func (t *SubAgentTool) Schema() json.RawMessage { + if len(t.spec.Schema) > 0 { + return t.spec.Schema + } + return subAgentSchema +} + +// ExecutionMode is parallel: independent sub-agents may run concurrently, since +// each spawns its own context and run (goroutine or process) with no shared +// mutable state. +func (t *SubAgentTool) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} + +// Execute spawns the child agent run and blocks until it settles, then returns +// the child's final assistant text as the tool result. The parent's ctx governs +// the child, so cancelling the parent run cancels in-flight sub-agents (in +// goroutine mode via ctx; in process mode via ctx cancelling the JSON-RPC call +// and Close killing the child). +func (t *SubAgentTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + // Goroutine mode requires NewRunConfig (it supplies the in-process provider + // stream). Process mode does not - the subprocess resolves its own provider + // from Process.Model - so the check is guarded to goroutine mode. This + // preserves the original precedence (nil NewRunConfig reported before an + // empty prompt) for the unchanged goroutine path. + if t.spec.Isolation != SubAgentIsolationProcess && t.spec.NewRunConfig == nil { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q: no run configuration", t.spec.Name) + } + var a subAgentArgs + if len(args) > 0 { + if err := json.Unmarshal(args, &a); err != nil { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q: decode args: %w", t.spec.Name, err) + } + } + if a.Prompt == "" { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q: empty prompt", t.spec.Name) + } + + if t.spec.Isolation == SubAgentIsolationProcess { + // Process mode returns only the child's final text (the JSON-RPC protocol + // does not stream partial updates), so onUpdate is intentionally not + // forwarded here; a caller supplying a sink gets no deltas in this mode. + return t.executeProcess(ctx, a.Prompt) + } + return t.executeGoroutine(ctx, id, a.Prompt, a.Description, onUpdate) +} + +// executeGoroutine runs the child agent loop in-process and returns its final +// text. This is the default mode and the original sub-agent behavior. +// +// id is the parent tool call's id and description is the (optional) task +// description; both are threaded onto any SubAgentProgressEvent emitted for this +// run so a consumer can key status by the parent task call. When the parent loop +// injected a run-level progress emitter into ctx (WithProgressEmitter), the +// child's tool-execution / turn boundaries are translated into +// SubAgentProgressEvent and surfaced up the parent stream; when no emitter is +// present (e.g. the tool is called directly in a unit test) progress reporting is +// silently skipped. +func (t *SubAgentTool) executeGoroutine(ctx context.Context, id, prompt, description string, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + // Concurrency guard: when a shared semaphore is configured, acquire a slot + // before spawning the child and release it via defer so a panic or error + // still frees the slot. A full channel blocks (queues) the acquire; a + // cancelled parent ctx abandons the wait instead of blocking forever. + if t.spec.Sem != nil { + select { + case t.spec.Sem <- struct{}{}: + defer func() { <-t.spec.Sem }() + case <-ctx.Done(): + return agentcore.AgentToolResult{}, ctx.Err() + } + } + runCfg := t.spec.NewRunConfig() + // Advertise the child's tools to the model. A spec may pin an explicit set + // (spec.Tools); otherwise fall back to the run config's registry — the tools + // the executor can actually run — so a factory that wires only the registry + // (like the generic task tool) still tells the child what it can call. + // Without this the model is handed an empty tool list, can only reply with + // text, and a delegated task that needs tools comes back empty. + tools := t.spec.Tools + if len(tools) == 0 && runCfg.Batch.ToolExecutorConfig.Registry != nil { + tools = runCfg.Batch.ToolExecutorConfig.Registry.List() + } + childCtx := &agentcore.AgentContext{ + SystemPrompt: t.spec.SystemPrompt, + Messages: agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent(prompt)}}, + }, + Tools: tools, + } + + stream := StartRun(ctx, childCtx, runCfg) + // Drain events (DrainStream never returns early, so the producer goroutine is + // never blocked on back-pressure); forward streamed child text as + // tool-execution updates when a sink is set. + var h StreamHandler + if onUpdate != nil { + h.OnText = func(delta string) { + onUpdate(agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(delta)}}) + } + } + // Progress reporting: when the parent loop injected a run-level emitter into + // ctx, translate the child's tool-execution / turn boundaries into + // SubAgentProgressEvent and emit them up the parent stream. Reporting is at + // activity granularity (per child tool start / turn boundary), NOT per text + // delta, so event volume stays proportional to the child's tool calls. When + // no emitter is present the OnEvent hook is left nil and progress is skipped. + if parentEmit := agentcore.ProgressEmitterFromContext(ctx); parentEmit != nil { + // chars accumulates the child's streamed text length so a coarse output + // token estimate can ride along on each progress event (0 = unknown). + chars := 0 + if prev := h.OnText; prev != nil { + h.OnText = func(delta string) { + chars += len(delta) + prev(delta) + } + } else { + h.OnText = func(delta string) { chars += len(delta) } + } + h.OnEvent = func(ev agentcore.AgentEvent) { + act := activityOf(ev) + if act == "" { + return + } + _ = parentEmit(ctx, agentcore.SubAgentProgressEvent{ + ToolCallID: id, + Description: description, + Activity: act, + Tokens: estimateTokens(chars), + }) + } + } + final, err := DrainStream(ctx, stream, h) + if err != nil { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q: %w", t.spec.Name, err) + } + text := "" + if final != nil { + text = agentcore.ContentToText(final.Content) + } + if text == "" { + text = fmt.Sprintf("(sub-agent %q produced no text output)", t.spec.Name) + } + // Surface a failed child run as a tool error so the parent model gets a + // signal the delegation failed (the tool executor marks the result + // IsError). A child whose final turn stopped on error/aborted otherwise + // looks like a successful delegation carrying error text. + if final != nil && (final.StopReason == agentcore.StopReasonError || final.StopReason == agentcore.StopReasonAborted) { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q failed (%s): %s", t.spec.Name, final.StopReason, text) + } + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(text)}}, nil +} + +// executeProcess runs the child agent loop in a fresh pigo subprocess over stdio +// JSON-RPC and returns its final text. A subprocess crash, transport error, or +// failed child run is surfaced as a tool error; the parent loop is unaffected. +// Streamed child text is not forwarded (the process protocol returns only the +// final result); the parent sees the complete result when the child settles. +func (t *SubAgentTool) executeProcess(ctx context.Context, prompt string) (agentcore.AgentToolResult, error) { + cfg := t.spec.Process + if cfg.Model == "" { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q: process mode requires Process.Model", t.spec.Name) + } + // Forward the child's tool names so the subprocess can rebuild a matching + // builtin set; an explicit ToolNames list wins over deriving from Tools. + toolNames := cfg.ToolNames + if len(toolNames) == 0 { + for _, tl := range t.spec.Tools { + toolNames = append(toolNames, tl.Name()) + } + } + params := SubAgentRunParams{ + Prompt: prompt, + SystemPrompt: t.spec.SystemPrompt, + Model: cfg.Model, + BaseURL: cfg.BaseURL, + Protocol: cfg.Protocol, + Tools: toolNames, + } + call := t.processCall + if call == nil { + call = defaultProcessCall + } + text, err := call(ctx, cfg, params) + if err != nil { + return agentcore.AgentToolResult{}, fmt.Errorf("sub-agent %q (process): %w", t.spec.Name, err) + } + if text == "" { + text = fmt.Sprintf("(sub-agent %q produced no text output)", t.spec.Name) + } + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(text)}}, nil +} + +// defaultProcessCall is the production subprocess transport: it launches the +// pigo binary (or cfg.Command) with the subagent-rpc flag, sends a single +// "subagent/run" JSON-RPC request over the child's stdin, and returns the +// child's final text from the response. The child is closed (killed if it does +// not exit on its own) before returning. A crash, transport error, or RPC error +// is returned as a Go error so executeProcess surfaces it as a tool error. +func defaultProcessCall(ctx context.Context, cfg SubAgentProcessConfig, params SubAgentRunParams) (string, error) { + command := cfg.Command + if command == "" { + exe, err := os.Executable() + if err != nil { + return "", fmt.Errorf("resolve pigo executable: %w", err) + } + command = exe + } + args := append([]string{SubAgentRPCFlag}, cfg.Args...) + client, err := jsonrpc.NewClient(jsonrpc.Config{ + Command: command, + Args: args, + Env: cfg.Env, + Dir: cfg.Dir, + Stderr: cfg.Stderr, + }) + if err != nil { + return "", err + } + defer client.Close() + raw, err := client.Call(ctx, SubAgentRPCMethod, params) + if err != nil { + return "", err + } + var res SubAgentRunResult + if err := json.Unmarshal(raw, &res); err != nil { + return "", fmt.Errorf("decode sub-agent result: %w", err) + } + return res.Text, nil +} + +// RunSubAgentOnce runs one sub-agent loop to completion and returns the child's +// final assistant text. It is the execution core shared by the process-isolated +// subprocess (cmd/pigo --subagent-rpc): given a resolved RunConfig (provider +// stream, tool registry) and the prompt/system prompt, it builds a fresh child +// context and drains the run. A run whose final turn stopped on error/aborted +// is reported as an error so the subprocess surfaces failure (as an RPC error) +// rather than returning empty text. It does not stream partial updates: the +// process protocol returns only the final result. +func RunSubAgentOnce(ctx context.Context, systemPrompt, prompt string, tools []agentcore.AgentTool, runCfg RunConfig) (string, error) { + childCtx := &agentcore.AgentContext{ + SystemPrompt: systemPrompt, + Messages: agentcore.MessageList{ + agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent(prompt)}}, + }, + Tools: tools, + } + stream := StartRun(ctx, childCtx, runCfg) + final, err := DrainStream(ctx, stream, StreamHandler{}) + if err != nil { + return "", err + } + text := "" + if final != nil { + text = agentcore.ContentToText(final.Content) + } + if final != nil && (final.StopReason == agentcore.StopReasonError || final.StopReason == agentcore.StopReasonAborted) { + // When the loop synthesizes an error turn (e.g. a provider connection + // failure) the diagnostic lands in ErrorMessage, not Content; fall back + // to it so the subprocess surfaces the real cause rather than a bare + // "error" stop reason. + if text == "" && final.ErrorMessage != "" { + text = final.ErrorMessage + } + if text == "" { + text = string(final.StopReason) + } + return text, fmt.Errorf("sub-agent failed (%s): %s", final.StopReason, text) + } + return text, nil +} diff --git a/pigo/internal/runtime/subagent_process_test.go b/pigo/internal/runtime/subagent_process_test.go new file mode 100644 index 0000000..9d4d605 --- /dev/null +++ b/pigo/internal/runtime/subagent_process_test.go @@ -0,0 +1,328 @@ +package runtime + +// Tests for process-isolated sub-agents (US-019, #135): the Isolation field, +// the process-mode Execute path (params shaping, crash-as-tool-error), the +// shared RunSubAgentOnce core, and the real defaultProcessCall transport +// (exec + stdio JSON-RPC + crash handling) driven through a tiny compiled +// helper binary - no provider or network, mirroring the plugin tests. + +import ( + "context" + "encoding/json" + "errors" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + + goruntime "runtime" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/provider" +) + +// errorTurn scripts a turn whose final message stops on StopReasonError, so a +// RunSubAgentOnce run surfaces as a failure (matching the goroutine-mode +// "failed run -> tool error" contract). +func errorTurn(msg string) fauxTurn { + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + final := partial + final.StopReason = agentcore.StopReasonError + final.ErrorMessage = msg + return fauxTurn{ + provider.StreamStartEvent{Partial: partial}, + provider.StreamDoneEvent{Message: final}, + } +} + +// TestSubAgentGoroutineModeUnchanged verifies the default (zero-value) isolation +// still runs in-process: the existing parent->child->parent tests cover the +// full path, but this pins that Isolation==goroutine is the default and reaches +// NewRunConfig (not the process path) so the regression is caught locally too. +func TestSubAgentGoroutineModeUnchanged(t *testing.T) { + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{textTurn("child answer")}, + } + sub := NewSubAgentTool(SubAgentSpec{ + Name: "researcher", + SystemPrompt: "you are a researcher", + NewRunConfig: func() RunConfig { return newFauxRunCfg(child) }, + }) + if sub.spec.Isolation != SubAgentIsolationGoroutine { + t.Errorf("default Isolation = %v, want goroutine", sub.spec.Isolation) + } + res, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil) + if err != nil { + t.Fatalf("goroutine Execute err = %v", err) + } + if got := agentcore.ContentToText(res.Content); got != "child answer" { + t.Errorf("goroutine result = %q, want 'child answer'", got) + } + if child.callCount() != 1 { + t.Errorf("child provider calls = %d, want 1", child.callCount()) + } +} + +// TestSubAgentGoroutineNilRunConfigErrors verifies goroutine mode reports a +// missing NewRunConfig, and - preserving the original precedence - reports it +// even when the prompt is empty (so a nil config is not masked by "empty +// prompt"). This pins the L3 regression: the NewRunConfig check stays before +// the empty-prompt check on the goroutine path. +func TestSubAgentGoroutineNilRunConfigErrors(t *testing.T) { + sub := NewSubAgentTool(SubAgentSpec{Name: "x"}) // no NewRunConfig, default goroutine + _, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil) + if err == nil || !strings.Contains(err.Error(), "no run configuration") { + t.Errorf("err = %v, want 'no run configuration'", err) + } + // Precedence: nil NewRunConfig is reported even when the prompt is empty. + _, err = sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":""}`), nil) + if err == nil || !strings.Contains(err.Error(), "no run configuration") { + t.Errorf("empty-prompt err = %v, want 'no run configuration' (nil config takes precedence)", err) + } +} + +// TestSubAgentProcessModeFake drives process mode through an injectable +// processCall (the test seam) to verify params shaping and crash-as-tool-error +// without a real subprocess. The real transport is covered separately by +// TestSubAgentProcessDefaultCall. +func TestSubAgentProcessModeFake(t *testing.T) { + t.Run("happy", func(t *testing.T) { + sub := NewSubAgentTool(SubAgentSpec{ + Name: "proc", + Isolation: SubAgentIsolationProcess, + Process: SubAgentProcessConfig{Model: "faux"}, + SystemPrompt: "you are a subprocess child", + }) + var got SubAgentRunParams + sub.processCall = func(ctx context.Context, cfg SubAgentProcessConfig, params SubAgentRunParams) (string, error) { + got = params + return "process result: 99", nil + } + res, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"find it"}`), nil) + if err != nil { + t.Fatalf("Execute err = %v", err) + } + if got := agentcore.ContentToText(res.Content); got != "process result: 99" { + t.Errorf("result = %q, want 'process result: 99'", got) + } + // The prompt, system prompt, and model are forwarded to the subprocess; + // NewRunConfig is NOT called (the subprocess resolves its own provider). + if got.Prompt != "find it" || got.Model != "faux" || got.SystemPrompt != "you are a subprocess child" { + t.Errorf("forwarded params = %+v", got) + } + }) + + t.Run("crash is tool error", func(t *testing.T) { + sub := NewSubAgentTool(SubAgentSpec{ + Name: "proc", + Isolation: SubAgentIsolationProcess, + Process: SubAgentProcessConfig{Model: "faux"}, + }) + sub.processCall = func(context.Context, SubAgentProcessConfig, SubAgentRunParams) (string, error) { + return "", errors.New("subprocess exited: signal: killed") + } + _, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"x"}`), nil) + if err == nil { + t.Fatal("expected tool error on subprocess crash, got nil") + } + }) + + t.Run("missing model errors", func(t *testing.T) { + sub := NewSubAgentTool(SubAgentSpec{ + Name: "proc", + Isolation: SubAgentIsolationProcess, + // Process.Model intentionally empty. + }) + _, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"x"}`), nil) + if err == nil { + t.Fatal("expected error when Process.Model is missing") + } + }) + + t.Run("forwards tool names", func(t *testing.T) { + // spec.Tools is in-process; process mode forwards only their names so + // the subprocess can rebuild builtins by name. + sub := NewSubAgentTool(SubAgentSpec{ + Name: "proc", + Isolation: SubAgentIsolationProcess, + Process: SubAgentProcessConfig{Model: "faux"}, + Tools: []agentcore.AgentTool{nameOnlyTool("read"), nameOnlyTool("grep")}, + }) + var got SubAgentRunParams + sub.processCall = func(_ context.Context, _ SubAgentProcessConfig, params SubAgentRunParams) (string, error) { + got = params + return "ok", nil + } + if _, err := sub.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"x"}`), nil); err != nil { + t.Fatalf("Execute err = %v", err) + } + if len(got.Tools) != 2 || got.Tools[0] != "read" || got.Tools[1] != "grep" { + t.Errorf("forwarded tool names = %v, want [read grep]", got.Tools) + } + }) +} + +// nameOnlyTool is a minimal AgentTool whose only meaningful attribute is its +// Name, used to verify process mode forwards tool names without needing real +// tool implementations. +func nameOnlyTool(name string) agentcore.AgentTool { return nameOnly{name: name} } + +type nameOnly struct{ name string } + +func (t nameOnly) Name() string { return t.name } +func (t nameOnly) Description() string { return "" } +func (t nameOnly) Schema() json.RawMessage { return json.RawMessage(`{}`) } +func (t nameOnly) ExecutionMode() agentcore.ToolExecutionMode { + return agentcore.ToolExecutionParallel +} +func (t nameOnly) Execute(context.Context, string, json.RawMessage, agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return agentcore.AgentToolResult{}, nil +} + +// TestRunSubAgentOnce verifies the shared subprocess-side agent core: a normal +// run returns the child's final text, and a run whose final turn stopped on +// error is reported as an error (so the subprocess surfaces it as an RPC error +// and the parent marks the tool result IsError). +func TestRunSubAgentOnce(t *testing.T) { + t.Run("success", func(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{textTurn("hello from child")}, + } + text, err := RunSubAgentOnce(context.Background(), "sys", "do it", nil, newFauxRunCfg(p)) + if err != nil { + t.Fatalf("err = %v", err) + } + if text != "hello from child" { + t.Errorf("text = %q, want 'hello from child'", text) + } + if p.callCount() != 1 { + t.Errorf("provider calls = %d, want 1", p.callCount()) + } + }) + + t.Run("failed run errors", func(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{errorTurn("boom")}, + } + _, err := RunSubAgentOnce(context.Background(), "sys", "do it", nil, newFauxRunCfg(p)) + if err == nil { + t.Fatal("expected error for failed child run, got nil") + } + // The diagnostic is in ErrorMessage (errorTurn sets no Content); the + // subprocess must surface it rather than a bare "error" stop reason. + if !strings.Contains(err.Error(), "boom") { + t.Errorf("error %q does not contain the 'boom' diagnostic", err.Error()) + } + }) +} + +// TestSubAgentProcessDefaultCall exercises the real defaultProcessCall transport +// (exec + stdio JSON-RPC) against a tiny compiled helper binary. It verifies the +// happy round-trip (prompt forwarded, result decoded) and that a crashing +// subprocess is surfaced as a Go error (the AC: "a subprocess crash is caught by the parent as a tool error"). +func TestSubAgentProcessDefaultCall(t *testing.T) { + bin := buildSubAgentHelper(t) + cfg := SubAgentProcessConfig{Command: bin} + + t.Run("happy round-trip", func(t *testing.T) { + text, err := defaultProcessCall(context.Background(), cfg, SubAgentRunParams{ + Prompt: "hello", Model: "faux", SystemPrompt: "sys", + }) + if err != nil { + t.Fatalf("defaultProcessCall err = %v", err) + } + if text != "echo: hello" { + t.Errorf("text = %q, want 'echo: hello'", text) + } + }) + + t.Run("crash is error", func(t *testing.T) { + _, err := defaultProcessCall(context.Background(), cfg, SubAgentRunParams{ + Prompt: "CRASH", Model: "faux", + }) + if err == nil { + t.Fatal("expected error for crashing subprocess, got nil") + } + }) +} + +// buildSubAgentHelper compiles the sub-agent helper binary into a temp dir and +// returns its path. The helper speaks the same JSON-RPC "subagent/run" wire +// format as pigo --subagent-rpc (decoded with plain encoding/json, no internal +// imports, so it stays a standalone main package). +func buildSubAgentHelper(t *testing.T) string { + t.Helper() + dir := t.TempDir() + srcPath := filepath.Join(dir, "subagent_helper.go") + if err := os.WriteFile(srcPath, []byte(subAgentHelperSrc), 0o644); err != nil { + t.Fatalf("write helper source: %v", err) + } + bin := filepath.Join(dir, "subagent_helper") + if goruntime.GOOS == "windows" { + bin += ".exe" + } + cmd := exec.Command("go", "build", "-o", bin, srcPath) + cmd.Env = os.Environ() + if out, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("build sub-agent helper: %v\n%s", err, out) + } + return bin +} + +// subAgentHelperSrc is a tiny JSON-RPC server mirroring pigo --subagent-rpc: +// read a "subagent/run" request per line, and either respond with +// {text:"echo: "} or, when the prompt is "CRASH", exit without +// responding to simulate a subprocess crash. +const subAgentHelperSrc = `package main + +import ( + "bufio" + "encoding/json" + "os" +) + +type params struct { + Prompt string ` + "`json:\"prompt\"`" + ` +} + +type request struct { + ID *json.RawMessage ` + "`json:\"id\"`" + ` + Method string ` + "`json:\"method\"`" + ` + Params params ` + "`json:\"params\"`" + ` +} + +func main() { + sc := bufio.NewScanner(os.Stdin) + sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + enc := json.NewEncoder(os.Stdout) + for sc.Scan() { + line := sc.Bytes() + if len(line) == 0 { + continue + } + var r request + if err := json.Unmarshal(line, &r); err != nil { + continue + } + if r.Params.Prompt == "CRASH" { + // Simulate a subprocess crash: exit without writing a response so + // the parent's JSON-RPC reader sees EOF and fails the call. + os.Exit(1) + } + result, _ := json.Marshal(map[string]string{"text": "echo: " + r.Params.Prompt}) + resp := map[string]any{ + "jsonrpc": "2.0", + "id": r.ID, + "result": json.RawMessage(result), + } + _ = enc.Encode(resp) + } +} +` diff --git a/pigo/internal/runtime/task.go b/pigo/internal/runtime/task.go new file mode 100644 index 0000000..7faef45 --- /dev/null +++ b/pigo/internal/runtime/task.go @@ -0,0 +1,141 @@ +// This file implements the generic `task` tool (US-002, #454): a general-purpose +// sub-agent the model can dispatch with a free-form prompt to fan out work in a +// single assistant message. It is a thin specialization of SubAgentTool - it +// reuses executeGoroutine's child-loop driving - configured with a generic +// system prompt (the delegated task comes from the call arguments at runtime), +// a shared concurrency semaphore, and a child tool set from which `task` itself +// is excluded (the nesting guard, wired in internal/cli/run). +package runtime + +import ( + "encoding/json" + "os" + "strconv" + "strings" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// DefaultMaxSubagents is the concurrency cap applied when PIGO_MAX_SUBAGENTS is +// unset or invalid. It bounds how many task sub-agents run at once so a fan-out +// cannot overwhelm the provider rate limit. +const DefaultMaxSubagents = 4 + +// taskDescription is advertised to the parent model so it knows when to delegate +// work to a generic sub-agent. +const taskDescription = "Dispatch a general-purpose sub-agent to autonomously complete a delegated task. " + + "The sub-agent runs its own agent loop with a fresh context and the standard tool set, then returns its final report. " + + "Provide a complete, self-contained prompt since the sub-agent shares none of this conversation's context. " + + "Multiple task calls in one message run in parallel." + +// taskSystemPrompt seeds every generic sub-agent's context. It is intentionally +// generic (the actual work arrives as the runtime prompt) and mirrors the +// parent agent's operating posture so a delegated task is carried out the same +// way the parent would. +const taskSystemPrompt = "You are a focused sub-agent working on one delegated task. " + + "You have your own fresh context and the standard tool set, but you cannot spawn further sub-agents. " + + "Complete the task fully using the tools available, then respond with a concise final report of what you did and any key findings. " + + "Your final message is returned verbatim to the agent that dispatched you, so make it self-contained." + +// taskSchema is the JSON Schema for a task invocation: a required self-contained +// prompt plus an optional short description used for status display. +var taskSchema = json.RawMessage(`{ + "type": "object", + "properties": { + "description": { + "type": "string", + "description": "A short (3-5 word) description of the task, for status display." + }, + "prompt": { + "type": "string", + "description": "The full task for the sub-agent to perform. It must be self-contained since the sub-agent runs with a fresh context and shares none of this conversation." + } + }, + "required": ["prompt"], + "additionalProperties": false +}`) + +// NewTaskTool builds the generic `task` sub-agent tool. factory produces a fresh +// child RunConfig per spawn (reusing the parent's provider stream/model and a +// child tool registry that must exclude `task` for the nesting guard); sem is a +// shared buffered channel bounding concurrent task runs (nil disables limiting). +// The child prompt comes from the call arguments at runtime, so a single generic +// spec serves every delegated task. +func NewTaskTool(factory func() RunConfig, sem chan struct{}) *SubAgentTool { + return NewSubAgentTool(SubAgentSpec{ + Name: "task", + Description: taskDescription, + SystemPrompt: taskSystemPrompt, + Schema: taskSchema, + NewRunConfig: factory, + Sem: sem, + }) +} + +// MaxSubagents resolves the concurrency cap for task sub-agents from +// PIGO_MAX_SUBAGENTS: absent or unparseable yields DefaultMaxSubagents (4), and +// a parsed value below 1 is floored to 1 so the semaphore always admits at least +// one runner. +func MaxSubagents() int { + v := strings.TrimSpace(os.Getenv("PIGO_MAX_SUBAGENTS")) + if v == "" { + return DefaultMaxSubagents + } + n, err := strconv.Atoi(v) + if err != nil { + return DefaultMaxSubagents + } + if n < 1 { + return 1 + } + return n +} + +// NewSubagentSemaphore builds the shared concurrency semaphore for task +// sub-agents, sized by MaxSubagents. One instance per run is created and shared +// across every task call so the cap is enforced run-wide. +func NewSubagentSemaphore() chan struct{} { + return make(chan struct{}, MaxSubagents()) +} + +// activityOf maps a child sub-agent event to the display verb surfaced in a +// SubAgentProgressEvent (D-8: tool name / phase, no argument summary). A child +// ToolExecutionStartEvent maps by tool name; a TurnStartEvent (a fresh turn with +// no tool in progress) maps to "Thinking". Every other event maps to "" so the +// caller emits nothing — progress is reported only at these activity boundaries, +// keeping event volume proportional to the child's tool calls rather than its +// text deltas (D-7). +func activityOf(ev agentcore.AgentEvent) string { + switch e := ev.(type) { + case agentcore.ToolExecutionStartEvent: + switch e.ToolName { + case "read": + return "Reading" + case "edit", "write": + return "Editing" + case "bash": + return "Running bash" + case "grep", "find", "ls": + return "Searching" + case "webfetch": + return "Fetching" + default: + return "" + } + case agentcore.TurnStartEvent: + return "Thinking" + default: + return "" + } +} + +// estimateTokens gives a coarse output-token estimate from a running character +// count of the child's streamed text (~4 chars per token). It rides along on +// each progress event as a rough "↓ tokens" figure; 0 means unknown (no text +// streamed yet). +func estimateTokens(chars int) int { + if chars <= 0 { + return 0 + } + return chars / 4 +} diff --git a/pigo/internal/runtime/task_test.go b/pigo/internal/runtime/task_test.go new file mode 100644 index 0000000..d1fab17 --- /dev/null +++ b/pigo/internal/runtime/task_test.go @@ -0,0 +1,233 @@ +package runtime + +// Tests for the generic task tool (US-002/003/004, #454): its identity/schema +// contract, the shared concurrency semaphore (N > cap never exceeds cap), the +// nesting guard (child tool set excludes "task"), that a task returns the +// child's final text, and that a failed child surfaces as a tool error. The +// child loop is driven through the faux provider seam (mirrors orchestration_test.go); +// only the provider boundary is faked. + +import ( + "context" + "encoding/json" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/agenttool" + "github.com/smallnest/pigo/internal/provider" +) + +// TestTaskToolContract pins the tool identity, parallel execution mode, and the +// {description?, prompt} schema with prompt required. +func TestTaskToolContract(t *testing.T) { + tool := NewTaskTool(func() RunConfig { return RunConfig{} }, nil) + if tool.Name() != "task" { + t.Errorf("Name() = %q, want task", tool.Name()) + } + if tool.ExecutionMode() != agentcore.ToolExecutionParallel { + t.Errorf("ExecutionMode() = %v, want parallel", tool.ExecutionMode()) + } + var schema struct { + Properties struct { + Description json.RawMessage `json:"description"` + Prompt json.RawMessage `json:"prompt"` + } `json:"properties"` + Required []string `json:"required"` + } + if err := json.Unmarshal(tool.Schema(), &schema); err != nil { + t.Fatalf("schema is not valid JSON: %v", err) + } + if len(schema.Properties.Prompt) == 0 || len(schema.Properties.Description) == 0 { + t.Errorf("schema must declare both prompt and description properties") + } + if len(schema.Required) != 1 || schema.Required[0] != "prompt" { + t.Errorf("required = %v, want [prompt]", schema.Required) + } +} + +// TestTaskReturnsChildText verifies a dispatched task drives an independent child +// loop and returns the child's final assistant text as the tool result. +func TestTaskReturnsChildText(t *testing.T) { + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{textTurn("child final report")}, + } + factory := func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: provider.StreamFnFromProvider(child)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: agenttool.NewToolRegistry()}}, + } + } + tool := NewTaskTool(factory, nil) + res, err := tool.Execute(context.Background(), "id", json.RawMessage(`{"description":"do x","prompt":"do the work"}`), nil) + if err != nil { + t.Fatalf("Execute err = %v", err) + } + if got := agentcore.ContentToText(res.Content); got != "child final report" { + t.Errorf("task result = %q, want 'child final report'", got) + } + if child.callCount() != 1 { + t.Errorf("child provider calls = %d, want 1", child.callCount()) + } +} + +// TestTaskFailedChildErrors verifies a child whose final turn stops on error is +// surfaced to the parent as a tool error (not a silent success). +func TestTaskFailedChildErrors(t *testing.T) { + // A child turn ending on StopReason=error, carrying diagnostic text as content + // (executeGoroutine surfaces the child's Content on failure). + errTurn := func(text string) fauxTurn { + partial := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant} + withText := partial + withText.Content = agentcore.ContentList{agentcore.NewTextContent(text)} + final := withText + final.StopReason = agentcore.StopReasonError + return fauxTurn{ + provider.StreamStartEvent{Partial: partial}, + provider.StreamTextEvent{Partial: withText}, + provider.StreamDoneEvent{Message: final}, + } + } + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{errTurn("child exploded")}, + } + factory := func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: provider.StreamFnFromProvider(child)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: agenttool.NewToolRegistry()}}, + } + } + tool := NewTaskTool(factory, nil) + _, err := tool.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil) + if err == nil { + t.Fatal("a child that stopped on error must surface as a tool error") + } + if !strings.Contains(err.Error(), "child exploded") { + t.Errorf("error should carry the child's diagnostic, got %v", err) + } +} + +// TestTaskSemaphoreBoundsConcurrency dispatches N tasks concurrently through a +// shared semaphore of capacity cap (< N) and asserts the number of children +// running at once never exceeds cap. Each child calls a blocking fake tool that +// parks on a barrier, so all admitted children pile up simultaneously and the +// peak concurrency is observable. +func TestTaskSemaphoreBoundsConcurrency(t *testing.T) { + const capN, n = 2, 6 + sem := make(chan struct{}, capN) + + var running, peak int64 + release := make(chan struct{}) + // blockTool parks until the test closes release, holding a semaphore slot for + // the duration and recording the peak number of concurrent children. + blockTool := execTool{ + name: "block", + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + cur := atomic.AddInt64(&running, 1) + for { + p := atomic.LoadInt64(&peak) + if cur <= p || atomic.CompareAndSwapInt64(&peak, p, cur) { + break + } + } + defer atomic.AddInt64(&running, -1) + select { + case <-release: + case <-ctx.Done(): + } + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent("blocked")}}, nil + }, + } + // Each child runs one turn that calls the blocking tool, then (after release) + // a final text turn. + factory := func() RunConfig { + p := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "c", ID: "c"}}, + turns: []fauxTurn{toolCallTurn("t", "block", `{}`), textTurn("done")}, + } + reg := agenttool.NewToolRegistry() + _ = reg.Register(blockTool) + return RunConfig{ + LoopConfig: LoopConfig{Model: "c", Stream: provider.StreamFnFromProvider(p)}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}}, + } + } + tool := NewTaskTool(factory, sem) + + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, _ = tool.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil) + }() + } + // Give the admitted children time to reach the barrier, then let them go. + deadline := time.After(2 * time.Second) + for atomic.LoadInt64(&running) < int64(capN) { + select { + case <-deadline: + t.Fatalf("only %d children started, expected the semaphore to admit %d", atomic.LoadInt64(&running), capN) + default: + time.Sleep(time.Millisecond) + } + } + // Hold briefly so any over-admission (a semaphore bug) would push peak > cap. + time.Sleep(50 * time.Millisecond) + close(release) + wg.Wait() + + if got := atomic.LoadInt64(&peak); got > int64(capN) { + t.Errorf("peak concurrent children = %d, must not exceed cap %d", got, capN) + } + if got := atomic.LoadInt64(&peak); got == 0 { + t.Error("no child ever ran; the semaphore blocked everything") + } +} + +// TestTaskAdvertisesRegistryTools verifies the child sub-agent is told about the +// tools it can actually run: when the spec pins no explicit tool set, the child +// context's Tools are populated from the run config's registry. Without this the +// model receives an empty tool list and cannot do real work (the "non-functional +// sub-agent" bug), so this guards the wiring, not just the result. +func TestTaskAdvertisesRegistryTools(t *testing.T) { + // Capture the tools the provider is handed for the child request. + var gotTools []agentcore.AgentTool + capturing := provider.StreamFn(func(ctx context.Context, model string, llm provider.LlmContext, cfg provider.StreamConfig) (*provider.AssistantMessageEventStream, error) { + gotTools = llm.Tools + child := &fauxProvider{ + name: "faux-child", + models: []provider.Model{{Provider: "faux-child", ID: "child"}}, + turns: []fauxTurn{textTurn("done")}, + } + return provider.StreamFnFromProvider(child)(ctx, model, llm, cfg) + }) + reg := agenttool.NewToolRegistry() + _ = reg.Register(echoTool("read", agentcore.ToolExecutionParallel, false)) + _ = reg.Register(echoTool("bash", agentcore.ToolExecutionParallel, false)) + factory := func() RunConfig { + return RunConfig{ + LoopConfig: LoopConfig{Model: "child", Stream: capturing}, + Batch: agenttool.BatchConfig{ToolExecutorConfig: agenttool.ToolExecutorConfig{Registry: reg}}, + } + } + tool := NewTaskTool(factory, nil) + if _, err := tool.Execute(context.Background(), "id", json.RawMessage(`{"prompt":"go"}`), nil); err != nil { + t.Fatalf("Execute err = %v", err) + } + if len(gotTools) != 2 { + t.Fatalf("child was advertised %d tools, want 2 (from the registry)", len(gotTools)) + } + names := map[string]bool{gotTools[0].Name(): true, gotTools[1].Name(): true} + if !names["read"] || !names["bash"] { + t.Errorf("child tools = %v, want read+bash from the registry", names) + } +} diff --git a/pigo/internal/runtime/telemetry.go b/pigo/internal/runtime/telemetry.go new file mode 100644 index 0000000..e238c92 --- /dev/null +++ b/pigo/internal/runtime/telemetry.go @@ -0,0 +1,137 @@ +// This file implements structured telemetry collection for a loop run +// (observability -- structured telemetry collection). A lightweight accumulator observes the AgentEvents the loop +// already emits and folds them into a compact summary — per-tool wall-clock +// durations, turn count, truncation count, compaction count, and the latest +// context-utilization ratio — that is surfaced once at run end as a +// TelemetryEvent (just before agent_end). +// +// The design is deliberately additive: telemetry rides the existing AgentEvent +// family and the existing stream-json output path, so no new dependency +// (Prometheus/OTLP) is introduced and existing stream-json consumers that do +// not know the "telemetry" event type keep working unchanged. Collection is +// passive — observing an event never changes loop behavior. +package runtime + +import ( + "sync" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// telemetry is the per-run accumulator. Most fields are folded from events on +// the single runLoop goroutine, but tool_execution_* events fire from the +// parallel tool-batch goroutines (ExecuteToolCalls), so all mutation goes +// through mu. +type telemetry struct { + mu sync.Mutex + turns int + truncationCount int + compactionCount int + + // toolStarts maps an in-flight tool call id to the wall-clock time its + // execution began, so the matching end event can compute a duration. Keying + // by call id (not tool name) keeps parallel tool batches correct. + toolStarts map[string]time.Time + // toolTimings aggregates finished tool durations by tool name. + toolTimings map[string]agentcore.ToolTiming + + // contextTokens / contextWindow capture the most recent context accounting so + // the summary can report the latest utilization ratio. contextWindow == 0 + // means the window is unknown (utilization is then reported as 0). + contextTokens int + contextWindow int + + // now is the clock, injectable for deterministic tests. Defaults to + // time.Now. + now func() time.Time +} + +// newTelemetry constructs an empty accumulator using the wall clock. +func newTelemetry() *telemetry { + return &telemetry{ + toolStarts: make(map[string]time.Time), + toolTimings: make(map[string]agentcore.ToolTiming), + now: time.Now, + } +} + +// observe folds a single emitted event into the accumulator. It is a no-op for +// event types that carry no telemetry signal, and it never mutates the event. +func (t *telemetry) observe(ev agentcore.AgentEvent) { + if t == nil { + return + } + t.mu.Lock() + defer t.mu.Unlock() + switch e := ev.(type) { + case agentcore.TurnStartEvent: + t.turns++ + case agentcore.ToolExecutionStartEvent: + t.toolStarts[e.ToolCallID] = t.now() + case agentcore.ToolExecutionEndEvent: + start, ok := t.toolStarts[e.ToolCallID] + if !ok { + return + } + delete(t.toolStarts, e.ToolCallID) + elapsed := t.now().Sub(start).Milliseconds() + if elapsed < 0 { + elapsed = 0 + } + agg := t.toolTimings[e.ToolName] + agg.Count++ + agg.TotalMs += elapsed + t.toolTimings[e.ToolName] = agg + case agentcore.TurnEndEvent: + if e.Message.StopReason == agentcore.StopReasonLength { + t.truncationCount++ + } + case agentcore.CompactionEvent: + // Count only successful compactions; a failed one (ErrorMessage set) left + // the context unchanged. + if e.ErrorMessage == "" { + t.compactionCount++ + } + } +} + +// recordContext captures the latest context-token usage and window so the +// summary can report the current utilization ratio. A non-positive window is +// treated as unknown. +func (t *telemetry) recordContext(tokens, window int) { + if t == nil { + return + } + t.mu.Lock() + defer t.mu.Unlock() + t.contextTokens = tokens + if window > 0 { + t.contextWindow = window + } +} + +// summary materializes the accumulated metrics into a TelemetryEvent. The +// per-tool map is copied so the emitted event does not alias the accumulator's +// live state. +func (t *telemetry) summary() agentcore.TelemetryEvent { + t.mu.Lock() + defer t.mu.Unlock() + timings := make(map[string]agentcore.ToolTiming, len(t.toolTimings)) + for name, v := range t.toolTimings { + timings[name] = v + } + var utilization float64 + if t.contextWindow > 0 { + utilization = float64(t.contextTokens) / float64(t.contextWindow) + } + return agentcore.TelemetryEvent{ + Turns: t.turns, + ToolDurationsMs: timings, + TruncationCount: t.truncationCount, + CompactionCount: t.compactionCount, + ContextUtilization: utilization, + ContextTokens: t.contextTokens, + ContextWindow: t.contextWindow, + } +} diff --git a/pigo/internal/runtime/telemetry_test.go b/pigo/internal/runtime/telemetry_test.go new file mode 100644 index 0000000..320fc86 --- /dev/null +++ b/pigo/internal/runtime/telemetry_test.go @@ -0,0 +1,276 @@ +package runtime + +// Tests for structured telemetry collection (observability -- structured telemetry collection, node-251): +// the loop accumulates per-tool durations, turn count, truncation count, +// compaction count, and the latest context-utilization ratio, then surfaces +// them as a TelemetryEvent emitted just before agent_end. Both the unit-level +// accumulator and the end-to-end loop wiring (including the stream-json +// headless surface a script reads) are covered. + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" + "github.com/smallnest/pigo/internal/compaction" + "github.com/smallnest/pigo/internal/provider" +) + +// fakeClock returns a now() closure that advances by step on every call, so a +// tool_execution_start/end pair yields a deterministic non-zero duration. +func fakeClock(start time.Time, step time.Duration) func() time.Time { + cur := start + return func() time.Time { + t := cur + cur = cur.Add(step) + return t + } +} + +// findTelemetry returns the first TelemetryEvent emitted, or nil. +func findTelemetry(events []agentcore.AgentEvent) *agentcore.TelemetryEvent { + for _, ev := range events { + if te, ok := ev.(agentcore.TelemetryEvent); ok { + return &te + } + } + return nil +} + +// TestTelemetryObserveToolTiming verifies a start/end pair records an aggregated +// per-tool duration, and repeated calls accumulate count and total. +func TestTelemetryObserveToolTiming(t *testing.T) { + tel := newTelemetry() + tel.now = fakeClock(time.Unix(0, 0), 5*time.Millisecond) + + // Two invocations of "echo": each start advances the clock 5ms, each end + // advances another 5ms, so each invocation measures 5ms. + for _, id := range []string{"c1", "c2"} { + tel.observe(agentcore.ToolExecutionStartEvent{ToolCallID: id, ToolName: "echo"}) + tel.observe(agentcore.ToolExecutionEndEvent{ToolCallID: id, ToolName: "echo"}) + } + + sum := tel.summary() + got, ok := sum.ToolDurationsMs["echo"] + if !ok { + t.Fatalf("expected timing for tool echo, got %+v", sum.ToolDurationsMs) + } + if got.Count != 2 { + t.Errorf("echo count = %d, want 2", got.Count) + } + if got.TotalMs != 10 { + t.Errorf("echo totalMs = %d, want 10", got.TotalMs) + } +} + +// TestTelemetryObserveCounters verifies turn/truncation/compaction counters and +// that a failed compaction is not counted. +func TestTelemetryObserveCounters(t *testing.T) { + tel := newTelemetry() + tel.observe(agentcore.TurnStartEvent{}) + tel.observe(agentcore.TurnStartEvent{}) + tel.observe(agentcore.TurnEndEvent{Message: agentcore.AssistantMessage{StopReason: agentcore.StopReasonLength}}) + tel.observe(agentcore.TurnEndEvent{Message: agentcore.AssistantMessage{StopReason: agentcore.StopReasonEndTurn}}) + tel.observe(agentcore.CompactionEvent{}) // success + tel.observe(agentcore.CompactionEvent{ErrorMessage: "boom"}) // failure, not counted + + sum := tel.summary() + if sum.Turns != 2 { + t.Errorf("turns = %d, want 2", sum.Turns) + } + if sum.TruncationCount != 1 { + t.Errorf("truncationCount = %d, want 1", sum.TruncationCount) + } + if sum.CompactionCount != 1 { + t.Errorf("compactionCount = %d, want 1 (failed compaction must not count)", sum.CompactionCount) + } +} + +// TestTelemetryContextUtilization verifies the ratio is used/window and is 0 +// when the window is unknown. +func TestTelemetryContextUtilization(t *testing.T) { + tel := newTelemetry() + tel.recordContext(500, 2000) + if got := tel.summary().ContextUtilization; got != 0.25 { + t.Errorf("utilization = %v, want 0.25", got) + } + + unknown := newTelemetry() + unknown.recordContext(500, 0) + if got := unknown.summary().ContextUtilization; got != 0 { + t.Errorf("utilization with unknown window = %v, want 0", got) + } +} + +// TestLoopEmitsTelemetryBeforeAgentEnd verifies the loop emits exactly one +// telemetry event immediately before agent_end and that it captures a tool +// execution. +func TestLoopEmitsTelemetryBeforeAgentEnd(t *testing.T) { + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + oneToolAssistant("c1", "echo"), + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("done")}}, + }), echoTool("echo", agentcore.ToolExecutionParallel, false)) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + + te := findTelemetry(events) + if te == nil { + t.Fatalf("expected a TelemetryEvent, got %+v", eventKinds(events)) + } + // Telemetry must be the penultimate event, immediately before agent_end. + if n := len(events); n < 2 || events[n-1].EventType() != agentcore.EventAgentEnd || events[n-2].EventType() != agentcore.EventTelemetry { + t.Errorf("telemetry must be emitted just before agent_end, got %v", eventKinds(events)) + } + if te.Turns != 2 { + t.Errorf("telemetry turns = %d, want 2", te.Turns) + } + if _, ok := te.ToolDurationsMs["echo"]; !ok { + t.Errorf("telemetry should record the echo tool timing, got %+v", te.ToolDurationsMs) + } + if te.ToolDurationsMs["echo"].Count != 1 { + t.Errorf("echo count = %d, want 1", te.ToolDurationsMs["echo"].Count) + } +} + +// TestLoopTelemetryCountsCompaction verifies a compaction that fires during the +// run increments the telemetry compaction counter and records a non-zero +// context-utilization ratio. +func TestLoopTelemetryCountsCompaction(t *testing.T) { + main := scriptedStream([]agentcore.AssistantMessage{ + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn, Content: agentcore.ContentList{agentcore.NewTextContent("ok")}}, + }) + cfg := newRunCfg(main) + cfg.SummaryStream = summaryStream("## Goal\ncompacted") + cfg.ContextWindow = 2000 + cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100} + agentCtx := &agentcore.AgentContext{Messages: bigUserMessages(12, 800)} + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + te := findTelemetry(events) + if te == nil { + t.Fatal("expected a TelemetryEvent") + } + if te.CompactionCount != 1 { + t.Errorf("telemetry compactionCount = %d, want 1", te.CompactionCount) + } + if te.ContextWindow != 2000 { + t.Errorf("telemetry contextWindow = %d, want 2000", te.ContextWindow) + } + if te.ContextUtilization <= 0 { + t.Errorf("telemetry contextUtilization = %v, want > 0", te.ContextUtilization) + } +} + +// TestLoopTelemetryCountsTruncation verifies a length-truncated assistant +// response increments the telemetry truncation counter. +func TestLoopTelemetryCountsTruncation(t *testing.T) { + truncated := oneToolAssistant("c1", "echo") + truncated.StopReason = agentcore.StopReasonLength + cfg := newRunCfg(scriptedStream([]agentcore.AssistantMessage{ + truncated, + {RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}, + }), echoTool("echo", agentcore.ToolExecutionParallel, false)) + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser}}} + + events := collectEvents(t, agentLoop(context.Background(), agentCtx, cfg)) + te := findTelemetry(events) + if te == nil { + t.Fatal("expected a TelemetryEvent") + } + if te.TruncationCount != 1 { + t.Errorf("telemetry truncationCount = %d, want 1", te.TruncationCount) + } +} + +// TestTelemetryEventEnvelope verifies the stream-json envelope carries every +// telemetry field, including a name-keyed per-tool timings object, so a script +// can read the metrics directly (headless surface, acceptance criterion 3). +func TestTelemetryEventEnvelope(t *testing.T) { + env := eventEnvelope(agentcore.TelemetryEvent{ + Turns: 3, + TruncationCount: 1, + CompactionCount: 2, + ContextUtilization: 0.5, + ContextTokens: 1000, + ContextWindow: 2000, + ToolDurationsMs: map[string]agentcore.ToolTiming{"echo": {Count: 2, TotalMs: 40}}, + }) + if env["type"] != agentcore.EventTelemetry { + t.Errorf("type = %v, want %q", env["type"], agentcore.EventTelemetry) + } + if env["turns"] != 3 || env["truncationCount"] != 1 || env["compactionCount"] != 2 { + t.Errorf("counter fields wrong: %+v", env) + } + if env["contextUtilization"] != 0.5 || env["contextTokens"] != 1000 || env["contextWindow"] != 2000 { + t.Errorf("context fields wrong: %+v", env) + } + tools, ok := env["toolDurationsMs"].(map[string]map[string]any) + if !ok { + t.Fatalf("toolDurationsMs type = %T, want map[string]map[string]any", env["toolDurationsMs"]) + } + if tools["echo"]["count"] != 2 || tools["echo"]["totalMs"] != int64(40) { + t.Errorf("echo timing wrong: %+v", tools["echo"]) + } +} + +// TestHeadlessStreamJSONSurfacesTelemetry verifies the run-end telemetry summary +// reaches the stream-json headless output as a parseable JSON line — the +// script-readable metric surface required by acceptance criterion 3. +func TestHeadlessStreamJSONSurfacesTelemetry(t *testing.T) { + p := &fauxProvider{ + name: "faux", + models: []provider.Model{{Provider: "faux", ID: "faux"}}, + turns: []fauxTurn{ + toolCallTurn("call-1", "echo", `{"msg":"hi"}`), + textTurn("done"), + }, + } + cfg := newFauxRunCfg(p, echoTool("echo", agentcore.ToolExecutionParallel, false)) + var out bytes.Buffer + agentCtx := &agentcore.AgentContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("start")}}}} + + if err := RunHeadless(context.Background(), agentCtx, HeadlessConfig{Run: cfg, Mode: StreamJSONMode, Out: &out}); err != nil { + t.Fatalf("RunHeadless stream-json: %v", err) + } + + var telemetryLine map[string]any + sc := bufio.NewScanner(&out) + for sc.Scan() { + line := bytes.TrimSpace(sc.Bytes()) + if len(line) == 0 { + continue + } + var env map[string]any + if err := json.Unmarshal(line, &env); err != nil { + t.Fatalf("stream-json line not valid JSON: %q (%v)", line, err) + } + if env["type"] == agentcore.EventTelemetry { + telemetryLine = env + } + } + if telemetryLine == nil { + t.Fatal("stream-json output must contain a telemetry event a script can read") + } + // turns is a JSON number; a script would read it as float64. + if turns, ok := telemetryLine["turns"].(float64); !ok || turns < 2 { + t.Errorf("telemetry turns = %v, want >= 2", telemetryLine["turns"]) + } + tools, ok := telemetryLine["toolDurationsMs"].(map[string]any) + if !ok || tools["echo"] == nil { + t.Errorf("telemetry must report the echo tool timing, got %v", telemetryLine["toolDurationsMs"]) + } +} + +// eventKinds maps events to their type strings for readable failure output. +func eventKinds(events []agentcore.AgentEvent) []string { + out := make([]string, len(events)) + for i, ev := range events { + out[i] = ev.EventType() + } + return out +} diff --git a/pigo/internal/runtime/template.go b/pigo/internal/runtime/template.go new file mode 100644 index 0000000..843f681 --- /dev/null +++ b/pigo/internal/runtime/template.go @@ -0,0 +1,182 @@ +// This file implements the prompt-template expansion engine (US-002, #332): it +// turns a template body plus tokenized invocation args into the final prompt +// text, supporting pi's positional/default/slice syntax. +// +// Supported placeholders (mirrors https://pi.dev/docs/latest/prompt-templates): +// - $1, $2, ... $N : Nth positional arg (1-indexed; out-of-range -> "") +// - $@, $ARGUMENTS : all args joined by a single space +// - ${1:-default} : arg 1 when present and non-empty, else `default` +// - ${@:-default}, ${ARGUMENTS:-default} : all args when non-empty, else default +// - ${@:N} : args from the Nth onward (1-indexed), joined +// - ${@:N:L} : L args starting at N, joined +// +// A single left-to-right pass expands both braced ${...} and bare $N/$@ forms. +// Because the pass consumes a ${...} as one unit, a bare $1 never matches inside +// ${1:-...}, and a default literal containing $ is not re-expanded (the +// substituted text is appended verbatim and the scan advances past it). +package runtime + +import ( + "strconv" + "strings" +) + +// ExpandTemplate expands template against the tokenized args. A template with no +// placeholder preserves ParseUserCommand's behavior: with no args it is returned +// verbatim, with args they are appended after a blank line (joined by spaces). +func ExpandTemplate(template string, args []string) string { + if !hasPlaceholder(template) { + if len(args) == 0 { + return template + } + return template + "\n\n" + strings.Join(args, " ") + } + var b strings.Builder + i := 0 + n := len(template) + for i < n { + c := template[i] + if c == '$' && i+1 < n { + next := template[i+1] + if next == '{' { + end := strings.IndexByte(template[i+2:], '}') + if end < 0 { + // No closing brace: emit the '$' literally and continue. + b.WriteByte(c) + i++ + continue + } + inner := template[i+2 : i+2+end] + b.WriteString(expandBraced(inner, args)) + i += 2 + end + 1 + continue + } + if next == '@' { + b.WriteString(strings.Join(args, " ")) + i += 2 + continue + } + if isDigitByte(next) { + j := i + 1 + for j < n && isDigitByte(template[j]) { + j++ + } + idx, _ := strconv.Atoi(template[i+1 : j]) + if idx >= 1 && idx <= len(args) { + b.WriteString(args[idx-1]) + } + i = j + continue + } + if strings.HasPrefix(template[i+1:], "ARGUMENTS") { + b.WriteString(strings.Join(args, " ")) + i += 1 + len("ARGUMENTS") + continue + } + } + b.WriteByte(c) + i++ + } + return b.String() +} + +// hasPlaceholder reports whether s contains any template placeholder: ${...}, +// $@, $, or $ARGUMENTS. A bare '$' not followed by one of these is not a +// placeholder (emitted literally), so a template like "price $5" still counts. +func hasPlaceholder(s string) bool { + for i := 0; i < len(s); i++ { + if s[i] != '$' { + continue + } + if i+1 >= len(s) { + return false + } + next := s[i+1] + if next == '{' || next == '@' || isDigitByte(next) { + return true + } + if strings.HasPrefix(s[i+1:], "ARGUMENTS") { + return true + } + } + return false +} + +// expandBraced expands the content inside ${...} (without the surrounding +// braces). It dispatches on the three forms: name:-default, @:N[:L], or plain. +func expandBraced(inner string, args []string) string { + // default form: name:-default + if idx := strings.Index(inner, ":-"); idx >= 0 { + return expandDefaulted(inner[:idx], inner[idx+2:], args) + } + // slice form: @:N or @:N:L + if idx := strings.Index(inner, ":"); idx >= 0 { + if name := inner[:idx]; name != "@" && name != "ARGUMENTS" { + return "" // slicing only applies to all-args + } + return expandSlice(inner[idx+1:], args) + } + // plain: N, @, or ARGUMENTS + return expandPlain(inner, args) +} + +// expandPlain expands a bare braced name (no :- or :): a positional index, @, +// or ARGUMENTS. An unknown name (e.g. ${foo}) expands to "". +func expandPlain(name string, args []string) string { + if name == "@" || name == "ARGUMENTS" { + return strings.Join(args, " ") + } + if idx, err := strconv.Atoi(name); err == nil { + if idx >= 1 && idx <= len(args) { + return args[idx-1] + } + return "" + } + return "" +} + +// expandDefaulted expands name:-default: the named positional (when in range and +// non-empty) or all-args (when the join is non-empty), otherwise the literal +// default. The default is not re-expanded. +func expandDefaulted(name, def string, args []string) string { + if name == "@" || name == "ARGUMENTS" { + if joined := strings.Join(args, " "); joined != "" { + return joined + } + return def + } + if idx, err := strconv.Atoi(name); err == nil { + if idx >= 1 && idx <= len(args) && args[idx-1] != "" { + return args[idx-1] + } + return def + } + return def +} + +// expandSlice expands the N (or N:L) part of ${@:N} / ${@:N:L}: args from the Nth +// onward (1-indexed), optionally limited to L, joined by spaces. N<1 or beyond +// the arg list yields "". +func expandSlice(rest string, args []string) string { + parts := strings.SplitN(rest, ":", 2) + start, _ := strconv.Atoi(parts[0]) + if start < 1 { + return "" + } + begin := start - 1 + if begin >= len(args) { + return "" + } + end := len(args) + if len(parts) == 2 { + if l, err := strconv.Atoi(parts[1]); err == nil && l >= 0 { + end = begin + l + if end > len(args) { + end = len(args) + } + } + } + return strings.Join(args[begin:end], " ") +} + +func isDigitByte(b byte) bool { return b >= '0' && b <= '9' } diff --git a/pigo/internal/runtime/template_test.go b/pigo/internal/runtime/template_test.go new file mode 100644 index 0000000..9856b88 --- /dev/null +++ b/pigo/internal/runtime/template_test.go @@ -0,0 +1,170 @@ +package runtime + +// Tests for the prompt-template expansion engine (US-002, #332). Covers the +// positional/default/slice syntax from the acceptance criteria, the +// ${...}-before-bare-$ ordering invariant, and the no-placeholder append +// behavior that preserves ParseUserCommand's existing semantics. + +import "testing" + +func TestExpandTemplateNoPlaceholder(t *testing.T) { + // No args: verbatim. + if got := ExpandTemplate("Take a note", nil); got != "Take a note" { + t.Errorf("no args: got %q", got) + } + if got := ExpandTemplate("Take a note", []string{}); got != "Take a note" { + t.Errorf("empty args: got %q", got) + } + // With args: appended after a blank line, joined by spaces. + if got := ExpandTemplate("Take a note", []string{"buy", "milk"}); got != "Take a note\n\nbuy milk" { + t.Errorf("with args: got %q", got) + } +} + +func TestExpandTemplatePositional(t *testing.T) { + args := []string{"Button", "click", "handler"} + if got := ExpandTemplate("$1", args); got != "Button" { + t.Errorf("$1 = %q, want Button", got) + } + if got := ExpandTemplate("$3", args); got != "handler" { + t.Errorf("$3 = %q, want handler", got) + } + // Out of range -> empty. + if got := ExpandTemplate("$5", args); got != "" { + t.Errorf("$5 = %q, want empty", got) + } + // Multi-digit. + if got := ExpandTemplate("$10", []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}); got != "j" { + t.Errorf("$10 = %q, want j", got) + } +} + +func TestExpandTemplateAllArgs(t *testing.T) { + args := []string{"Button", "click", "handler"} + if got := ExpandTemplate("$@", args); got != "Button click handler" { + t.Errorf("$@ = %q", got) + } + if got := ExpandTemplate("$ARGUMENTS", args); got != "Button click handler" { + t.Errorf("$ARGUMENTS = %q", got) + } +} + +func TestExpandTemplateDefaultPositional(t *testing.T) { + // Present and non-empty -> the arg. + if got := ExpandTemplate("${1:-default}", []string{"a"}); got != "a" { + t.Errorf("present: got %q", got) + } + // Absent -> default. + if got := ExpandTemplate("${1:-default}", nil); got != "default" { + t.Errorf("absent: got %q", got) + } + // Present but empty -> default. + if got := ExpandTemplate("${1:-default}", []string{""}); got != "default" { + t.Errorf("empty: got %q", got) + } + // Default containing a $ is not re-expanded. + if got := ExpandTemplate("${1:-pay $5}", nil); got != "pay $5" { + t.Errorf("default literal $: got %q", got) + } +} + +func TestExpandTemplateDefaultAllArgs(t *testing.T) { + if got := ExpandTemplate("${@:-default}", []string{"a", "b"}); got != "a b" { + t.Errorf("present: got %q", got) + } + if got := ExpandTemplate("${@:-default}", nil); got != "default" { + t.Errorf("absent: got %q", got) + } + if got := ExpandTemplate("${ARGUMENTS:-default}", []string{"x"}); got != "x" { + t.Errorf("ARGUMENTS present: got %q", got) + } + if got := ExpandTemplate("${ARGUMENTS:-default}", nil); got != "default" { + t.Errorf("ARGUMENTS absent: got %q", got) + } +} + +func TestExpandTemplateSlice(t *testing.T) { + args := []string{"a", "b", "c", "d", "e"} + // ${@:N}: from Nth onward. + if got := ExpandTemplate("${@:2}", args); got != "b c d e" { + t.Errorf("${@:2} = %q", got) + } + if got := ExpandTemplate("${@:1}", args); got != "a b c d e" { + t.Errorf("${@:1} = %q", got) + } + // ${@:N:L}: L args starting at N. + if got := ExpandTemplate("${@:2:2}", args); got != "b c" { + t.Errorf("${@:2:2} = %q", got) + } + if got := ExpandTemplate("${@:1:3}", args); got != "a b c" { + t.Errorf("${@:1:3} = %q", got) + } + // N beyond args -> empty. + if got := ExpandTemplate("${@:9}", args); got != "" { + t.Errorf("${@:9} = %q, want empty", got) + } + // L clamps to available. + if got := ExpandTemplate("${@:3:100}", args); got != "c d e" { + t.Errorf("${@:3:100} = %q, want c d e", got) + } + // N<1 -> empty. + if got := ExpandTemplate("${@:0}", args); got != "" { + t.Errorf("${@:0} = %q, want empty", got) + } +} + +func TestExpandTemplateOrderingNoDoubleMatch(t *testing.T) { + // A bare $1 must NOT match inside ${1:-...}. With args present, ${1:-$2} + // uses arg1 ("a"); the $2 default is never consulted nor expanded. + if got := ExpandTemplate("${1:-$2}", []string{"a"}); got != "a" { + t.Errorf("${1:-$2} with arg1=a: got %q, want a", got) + } + // With arg1 absent, the default "$2" is taken literally (not expanded to ""). + if got := ExpandTemplate("${1:-$2}", nil); got != "$2" { + t.Errorf("${1:-$2} absent: got %q, want literal $2", got) + } + // Braced and bare coexist in one pass. + if got := ExpandTemplate("${1:-x} and $2", []string{"a", "b"}); got != "a and b" { + t.Errorf("coexist: got %q", got) + } +} + +func TestExpandTemplateExampleFromSpec(t *testing.T) { + // $@ expands to ALL args joined (AC: "all args joined by a single space"). + if got := ExpandTemplate("echo $@", []string{"a", "b", "c"}); got != "echo a b c" { + t.Errorf("$@ all-args: got %q", got) + } + // The pi "component" example: $1 for the name, features are the remaining + // args expressed with the slice form ${@:2} (the right tool for "rest"). + tmpl := "Create a React component named $1 with features: ${@:2}" + if got := ExpandTemplate(tmpl, []string{"Button", "click", "handler"}); got != "Create a React component named Button with features: click handler" { + t.Errorf("component example: got %q", got) + } + // "Summarize the current state in ${1:-7} bullet points." + bullets := "Summarize the current state in ${1:-7} bullet points." + if got := ExpandTemplate(bullets, nil); got != "Summarize the current state in 7 bullet points." { + t.Errorf("default bullets: got %q", got) + } + if got := ExpandTemplate(bullets, []string{"5"}); got != "Summarize the current state in 5 bullet points." { + t.Errorf("explicit bullets: got %q", got) + } +} + +func TestExpandTemplateLiteralDollar(t *testing.T) { + // A '$' not forming a placeholder is emitted literally. Such a template has + // no placeholder, so args (if any) are appended per the no-placeholder rule. + if got := ExpandTemplate("100$ off", nil); got != "100$ off" { + t.Errorf("literal $ mid-string: got %q", got) + } + if got := ExpandTemplate("trailing$", nil); got != "trailing$" { + t.Errorf("trailing $ no args: got %q", got) + } + // No placeholder + args -> append after a blank line. + if got := ExpandTemplate("trailing$", []string{"x"}); got != "trailing$\n\nx" { + t.Errorf("trailing $ with args: got %q", got) + } + // A literal '$' alongside a real placeholder: '$' before a space stays '$'. + if got := ExpandTemplate("cost $ and $1", []string{"five"}); got != "cost $ and five" { + t.Errorf("literal $ next to placeholder: got %q", got) + } +} diff --git a/pigo/internal/runtime/testtools_test.go b/pigo/internal/runtime/testtools_test.go new file mode 100644 index 0000000..8e36e60 --- /dev/null +++ b/pigo/internal/runtime/testtools_test.go @@ -0,0 +1,52 @@ +package runtime + +import ( + "context" + "encoding/json" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// execTool is a configurable AgentTool used by the loop/headless tests that +// remain in package agent. Its canonical definition moved to +// internal/agenttool with tool_executor_test.go (US-003 of the package split); +// this copy is re-provided here so the agent-resident tests keep compiling +// during the transition. +type execTool struct { + name string + schema string + run func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) + mode agentcore.ToolExecutionMode +} + +func (t execTool) Name() string { return t.name } +func (t execTool) Description() string { return "exec" } +func (t execTool) Schema() json.RawMessage { + if t.schema == "" { + return nil + } + return json.RawMessage(t.schema) +} +func (t execTool) ExecutionMode() agentcore.ToolExecutionMode { + if t.mode == "" { + return agentcore.ToolExecutionParallel + } + return t.mode +} +func (t execTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + return t.run(ctx, id, args, onUpdate) +} + +// echoTool returns its name as text; optionally terminates. Canonical +// definition moved with batch_executor_test.go; re-provided here for the +// agent-resident tests. +func echoTool(name string, mode agentcore.ToolExecutionMode, terminate bool) execTool { + return execTool{ + name: name, + mode: mode, + run: func(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) { + term := terminate + return agentcore.AgentToolResult{Content: agentcore.ContentList{agentcore.NewTextContent(name)}, Terminate: &term}, nil + }, + } +} diff --git a/pigo/internal/selfupdate/cache.go b/pigo/internal/selfupdate/cache.go new file mode 100644 index 0000000..77cb26a --- /dev/null +++ b/pigo/internal/selfupdate/cache.go @@ -0,0 +1,102 @@ +// This file caches the latest-release check so pigo's startup banner can show +// "update available" without a network call on every launch (US-004, FR-10). +// The cache lives at $PIGO_HOME/update-check.json (or ~/.pigo/update-check.json) +// and records the last check time plus the latest tag seen. CachedLatest reads +// it synchronously (fast, local); StartBackgroundCheck refreshes it off the hot +// path when older than the TTL, so a fresh result shows on the next launch. All +// failures are silent: a missing or corrupt cache, an unresolvable home, or a +// network error never surfaces an error or blocks startup. +package selfupdate + +import ( + "context" + "encoding/json" + "net/http" + "os" + "path/filepath" + "time" +) + +// checkTTL is the minimum interval between networked latest-release checks. +const checkTTL = 24 * time.Hour + +// cacheFileName is the on-disk cache under the pigo home directory. +const cacheFileName = "update-check.json" + +// updateCache is the on-disk shape of the latest-release check cache. +type updateCache struct { + CheckedAt time.Time `json:"checked_at"` + Latest string `json:"latest"` +} + +// cachePath returns the cache file path, or "" when the home dir is unavailable. +func cachePath() string { + if dir := os.Getenv("PIGO_HOME"); dir != "" { + return filepath.Join(dir, cacheFileName) + } + home, err := os.UserHomeDir() + if err != nil { + return "" + } + return filepath.Join(home, ".pigo", cacheFileName) +} + +// CachedLatest returns the latest tag recorded in the cache and whether the cache +// is still fresh (younger than checkTTL). A missing, corrupt, or unreadable cache +// yields ("", false). It never returns an error — the banner must not break on a +// bad cache. +func CachedLatest() (latest string, fresh bool) { + p := cachePath() + if p == "" { + return "", false + } + data, err := os.ReadFile(p) + if err != nil { + return "", false + } + var c updateCache + if err := json.Unmarshal(data, &c); err != nil { + return "", false + } + return c.Latest, time.Since(c.CheckedAt) < checkTTL +} + +// StartBackgroundCheck refreshes the cache off the hot path when it is stale +// (older than checkTTL). It returns immediately; the actual network check runs in +// a goroutine so it never blocks banner rendering or first input. When current is +// not a release version (dev/unknown) it does nothing — there is nothing to +// compare against. All errors are swallowed: a failed check simply leaves the +// cache untouched for next time. +func StartBackgroundCheck(current string) { + if !IsReleaseVersion(current) { + return + } + if _, fresh := CachedLatest(); fresh { + return + } + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + tag, err := LatestTag(ctx, &http.Client{Timeout: 10 * time.Second}, Repo) + if err != nil || tag == "" { + return + } + writeCache(tag) + }() +} + +// writeCache persists the latest tag with the current time. Failures are silent. +func writeCache(latest string) { + p := cachePath() + if p == "" { + return + } + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + return + } + data, err := json.Marshal(updateCache{CheckedAt: time.Now(), Latest: latest}) + if err != nil { + return + } + _ = os.WriteFile(p, data, 0o644) +} diff --git a/pigo/internal/selfupdate/cache_test.go b/pigo/internal/selfupdate/cache_test.go new file mode 100644 index 0000000..727fa98 --- /dev/null +++ b/pigo/internal/selfupdate/cache_test.go @@ -0,0 +1,52 @@ +package selfupdate + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + "time" +) + +func TestCachedLatest(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_HOME", dir) + + // No cache file yet → not fresh, empty latest. + if latest, fresh := CachedLatest(); latest != "" || fresh { + t.Errorf("empty cache = (%q,%v), want (\"\",false)", latest, fresh) + } + + // Fresh cache → returns latest and fresh=true. + writeCache("v0.4.0") + if latest, fresh := CachedLatest(); latest != "v0.4.0" || !fresh { + t.Errorf("fresh cache = (%q,%v), want (v0.4.0,true)", latest, fresh) + } + + // Stale cache (older than TTL) → latest kept, fresh=false. + stale, _ := json.Marshal(updateCache{CheckedAt: time.Now().Add(-25 * time.Hour), Latest: "v0.3.0"}) + if err := os.WriteFile(filepath.Join(dir, cacheFileName), stale, 0o644); err != nil { + t.Fatal(err) + } + if latest, fresh := CachedLatest(); latest != "v0.3.0" || fresh { + t.Errorf("stale cache = (%q,%v), want (v0.3.0,false)", latest, fresh) + } + + // Corrupt cache → silent ("", false), never an error. + if err := os.WriteFile(filepath.Join(dir, cacheFileName), []byte("{not json"), 0o644); err != nil { + t.Fatal(err) + } + if latest, fresh := CachedLatest(); latest != "" || fresh { + t.Errorf("corrupt cache = (%q,%v), want (\"\",false)", latest, fresh) + } +} + +func TestStartBackgroundCheckDevNoWrite(t *testing.T) { + dir := t.TempDir() + t.Setenv("PIGO_HOME", dir) + // dev is not a release version → must not touch the network or write a cache. + StartBackgroundCheck("dev") + if _, err := os.Stat(filepath.Join(dir, cacheFileName)); !os.IsNotExist(err) { + t.Errorf("dev build wrote a cache file; want none") + } +} diff --git a/pigo/internal/selfupdate/update.go b/pigo/internal/selfupdate/update.go new file mode 100644 index 0000000..f9e7611 --- /dev/null +++ b/pigo/internal/selfupdate/update.go @@ -0,0 +1,264 @@ +// This file implements pigo's binary self-replacement for `pigo update` (issue +// #466). Given the current build version it discovers the latest release (via +// version.go), downloads the matching goreleaser archive for the running +// GOOS/GOARCH, verifies its SHA256 against the release's checksums.txt, and +// atomically replaces the running executable. +// +// The archive naming mirrors .goreleaser.yaml and install.sh exactly, so this +// stays a single source of truth with the release tooling. Replacement is +// atomic: the new binary is written to a temp file in the target's directory +// and os.Rename'd over the current executable, so a failure mid-download never +// leaves a truncated binary in place. +package selfupdate + +import ( + "archive/tar" + "archive/zip" + "bufio" + "bytes" + "compress/gzip" + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "runtime" + "strings" + "time" +) + +// checksumsFile is the goreleaser checksums artifact name (see .goreleaser.yaml). +const checksumsFile = "checksums.txt" + +// Updater performs a self-replacement. Its fields are seams for testing; +// NewUpdater fills them with production defaults. +type Updater struct { + HTTPClient *http.Client + // Repo is "owner/name"; ReleaseBaseURL overrides the download host in tests. + Repo string + ReleaseBaseURL string // e.g. https://github.com/smallnest/pigo/releases/download + GOOS, GOARCH string + // ExecPath is the executable to replace; defaults to os.Executable(). + ExecPath string +} + +// Run performs `pigo update` (self-update pigo). It discovers the latest +// release, compares it to current, and replaces the running binary when a +// newer release exists. When current is a source build ("dev"), it cannot +// compare and proceeds to install the latest. Returns a process exit code. +func Run(ctx context.Context, current string, out, errOut io.Writer) int { + tag, err := LatestTag(ctx, nil, Repo) + if err != nil { + fmt.Fprintf(errOut, "pigo: failed to check for updates: %v\n", err) + return 1 + } + if avail, comparable := UpdateAvailable(current, tag); comparable && !avail { + fmt.Fprintf(out, "already up to date at %s\n", current) + return 0 + } + u, err := NewUpdater() + if err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + if err := u.Apply(ctx, tag, out); err != nil { + fmt.Fprintf(errOut, "pigo: %v\n", err) + return 1 + } + fmt.Fprintf(out, "updated to %s\n", tag) + return 0 +} + +// NewUpdater returns an Updater configured for the running process. +func NewUpdater() (*Updater, error) { + exe, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("selfupdate: locate executable: %w", err) + } + // Resolve symlinks so we replace the real file, not a symlink. + if resolved, err := filepath.EvalSymlinks(exe); err == nil { + exe = resolved + } + return &Updater{ + HTTPClient: &http.Client{Timeout: 60 * time.Second}, + Repo: Repo, + ReleaseBaseURL: "https://github.com/" + Repo + "/releases/download", + GOOS: runtime.GOOS, + GOARCH: runtime.GOARCH, + ExecPath: exe, + }, nil +} + +// archiveName builds the goreleaser archive filename for a release version +// (without leading "v") on the updater's platform. It mirrors the +// name_template and format_overrides in .goreleaser.yaml. +func (u *Updater) archiveName(versionNoV string) string { + osName := map[string]string{"darwin": "Darwin", "linux": "Linux", "windows": "Windows"}[u.GOOS] + if osName == "" { + osName = u.GOOS + } + arch := map[string]string{"amd64": "x86_64", "386": "i386"}[u.GOARCH] + if arch == "" { + arch = u.GOARCH // arm64 and others pass through + } + ext := "tar.gz" + if u.GOOS == "windows" { + ext = "zip" + } + return fmt.Sprintf("pigo_%s_%s_%s.%s", versionNoV, osName, arch, ext) +} + +// binaryName is the executable name inside the archive. +func (u *Updater) binaryName() string { + if u.GOOS == "windows" { + return "pigo.exe" + } + return "pigo" +} + +// Apply downloads the release identified by tag, verifies its checksum, and +// atomically replaces the target executable. tag is like "v0.4.0". +func (u *Updater) Apply(ctx context.Context, tag string, out io.Writer) error { + versionNoV := strings.TrimPrefix(strings.TrimSpace(tag), "v") + archive := u.archiveName(versionNoV) + base := fmt.Sprintf("%s/%s", strings.TrimRight(u.ReleaseBaseURL, "/"), tag) + + fmt.Fprintf(out, "downloading %s ...\n", archive) + archiveBytes, err := u.download(ctx, base+"/"+archive) + if err != nil { + return fmt.Errorf("selfupdate: download archive: %w", err) + } + + sums, err := u.download(ctx, base+"/"+checksumsFile) + if err != nil { + return fmt.Errorf("selfupdate: download checksums: %w", err) + } + want, err := checksumFor(sums, archive) + if err != nil { + return err + } + got := sha256.Sum256(archiveBytes) + if hex.EncodeToString(got[:]) != want { + return fmt.Errorf("selfupdate: checksum mismatch for %s (archive corrupt or tampered)", archive) + } + + binary, err := extractBinary(archiveBytes, u.binaryName(), u.GOOS == "windows") + if err != nil { + return err + } + if err := u.replace(binary); err != nil { + return err + } + return nil +} + +// download fetches url and returns the full body. A non-200 status is an error. +func (u *Updater) download(ctx context.Context, url string) ([]byte, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + resp, err := u.HTTPClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("GET %s: %s", url, resp.Status) + } + return io.ReadAll(resp.Body) +} + +// replace atomically swaps the target executable with newBin: it writes a temp +// file in the target's directory, sets it executable, and renames it over the +// target. Writing to the same directory keeps the rename atomic (same +// filesystem). A permission error on the directory yields an actionable message. +func (u *Updater) replace(newBin []byte) error { + dir := filepath.Dir(u.ExecPath) + tmp, err := os.CreateTemp(dir, ".pigo-update-*") + if err != nil { + return fmt.Errorf("selfupdate: cannot write to %s: %w (try running with sudo, or install pigo to a writable directory)", dir, err) + } + tmpName := tmp.Name() + defer os.Remove(tmpName) // no-op after successful rename + + if _, err := tmp.Write(newBin); err != nil { + tmp.Close() + return fmt.Errorf("selfupdate: write new binary: %w", err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("selfupdate: close new binary: %w", err) + } + if err := os.Chmod(tmpName, 0o755); err != nil { + return fmt.Errorf("selfupdate: chmod new binary: %w", err) + } + if err := os.Rename(tmpName, u.ExecPath); err != nil { + return fmt.Errorf("selfupdate: replace %s: %w", u.ExecPath, err) + } + return nil +} + +// checksumFor finds the hex SHA256 for archive in a goreleaser checksums.txt +// body (lines of " "). +func checksumFor(sums []byte, archive string) (string, error) { + sc := bufio.NewScanner(bytes.NewReader(sums)) + for sc.Scan() { + fields := strings.Fields(sc.Text()) + if len(fields) == 2 && fields[1] == archive { + return fields[0], nil + } + } + return "", fmt.Errorf("selfupdate: %s not found in checksums.txt", archive) +} + +// extractBinary pulls the named binary out of an archive (tar.gz, or zip when +// isZip). It returns the binary bytes. +func extractBinary(archive []byte, name string, isZip bool) ([]byte, error) { + if isZip { + return extractFromZip(archive, name) + } + return extractFromTarGz(archive, name) +} + +func extractFromTarGz(archive []byte, name string) ([]byte, error) { + gz, err := gzip.NewReader(bytes.NewReader(archive)) + if err != nil { + return nil, fmt.Errorf("selfupdate: gzip reader: %w", err) + } + defer gz.Close() + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + if err != nil { + return nil, fmt.Errorf("selfupdate: read tar: %w", err) + } + if filepath.Base(hdr.Name) == name && hdr.Typeflag == tar.TypeReg { + return io.ReadAll(tr) + } + } + return nil, fmt.Errorf("selfupdate: %s not found in archive", name) +} + +func extractFromZip(archive []byte, name string) ([]byte, error) { + zr, err := zip.NewReader(bytes.NewReader(archive), int64(len(archive))) + if err != nil { + return nil, fmt.Errorf("selfupdate: zip reader: %w", err) + } + for _, f := range zr.File { + if filepath.Base(f.Name) == name { + rc, err := f.Open() + if err != nil { + return nil, fmt.Errorf("selfupdate: open %s in zip: %w", name, err) + } + defer rc.Close() + return io.ReadAll(rc) + } + } + return nil, fmt.Errorf("selfupdate: %s not found in archive", name) +} diff --git a/pigo/internal/selfupdate/update_test.go b/pigo/internal/selfupdate/update_test.go new file mode 100644 index 0000000..3475e55 --- /dev/null +++ b/pigo/internal/selfupdate/update_test.go @@ -0,0 +1,175 @@ +package selfupdate + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" +) + +func TestArchiveName(t *testing.T) { + tests := []struct { + goos, goarch, want string + }{ + {"darwin", "arm64", "pigo_0.4.0_Darwin_arm64.tar.gz"}, + {"darwin", "amd64", "pigo_0.4.0_Darwin_x86_64.tar.gz"}, + {"linux", "amd64", "pigo_0.4.0_Linux_x86_64.tar.gz"}, + {"linux", "386", "pigo_0.4.0_Linux_i386.tar.gz"}, + {"windows", "amd64", "pigo_0.4.0_Windows_x86_64.zip"}, + } + for _, tt := range tests { + u := &Updater{GOOS: tt.goos, GOARCH: tt.goarch} + if got := u.archiveName("0.4.0"); got != tt.want { + t.Errorf("archiveName(%s/%s) = %q, want %q", tt.goos, tt.goarch, got, tt.want) + } + } +} + +func TestChecksumFor(t *testing.T) { + sums := []byte("abc123 pigo_0.4.0_Linux_x86_64.tar.gz\ndef456 pigo_0.4.0_Darwin_arm64.tar.gz\n") + got, err := checksumFor(sums, "pigo_0.4.0_Darwin_arm64.tar.gz") + if err != nil || got != "def456" { + t.Errorf("checksumFor = (%q,%v), want (def456,nil)", got, err) + } + if _, err := checksumFor(sums, "missing.tar.gz"); err == nil { + t.Error("expected error for missing archive") + } +} + +func TestExtractBinaryTarGz(t *testing.T) { + want := []byte("#!fake pigo binary") + archive := makeTarGz(t, "pigo", want) + got, err := extractBinary(archive, "pigo", false) + if err != nil { + t.Fatalf("extractBinary: %v", err) + } + if !bytes.Equal(got, want) { + t.Errorf("extracted = %q, want %q", got, want) + } + if _, err := extractBinary(archive, "nope", false); err == nil { + t.Error("expected error for missing binary") + } +} + +func TestReplaceAtomic(t *testing.T) { + dir := t.TempDir() + target := filepath.Join(dir, "pigo") + if err := os.WriteFile(target, []byte("old"), 0o755); err != nil { + t.Fatal(err) + } + u := &Updater{ExecPath: target} + newBin := []byte("new binary content") + if err := u.replace(newBin); err != nil { + t.Fatalf("replace: %v", err) + } + got, _ := os.ReadFile(target) + if !bytes.Equal(got, newBin) { + t.Errorf("after replace = %q, want %q", got, newBin) + } + // No leftover temp files in the directory. + entries, _ := os.ReadDir(dir) + if len(entries) != 1 { + t.Errorf("expected 1 file after replace, got %d", len(entries)) + } +} + +func TestApplyEndToEnd(t *testing.T) { + binary := []byte("brand new pigo v0.4.0") + archive := makeTarGz(t, "pigo", binary) + sum := sha256.Sum256(archive) + archiveName := "pigo_0.4.0_Linux_x86_64.tar.gz" + sums := fmt.Sprintf("%s %s\n", hex.EncodeToString(sum[:]), archiveName) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch filepath.Base(r.URL.Path) { + case archiveName: + _, _ = w.Write(archive) + case checksumsFile: + _, _ = w.Write([]byte(sums)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + dir := t.TempDir() + target := filepath.Join(dir, "pigo") + _ = os.WriteFile(target, []byte("old"), 0o755) + + u := &Updater{ + HTTPClient: srv.Client(), + Repo: "smallnest/pigo", + ReleaseBaseURL: srv.URL, + GOOS: "linux", + GOARCH: "amd64", + ExecPath: target, + } + if err := u.Apply(context.Background(), "v0.4.0", &bytes.Buffer{}); err != nil { + t.Fatalf("Apply: %v", err) + } + got, _ := os.ReadFile(target) + if !bytes.Equal(got, binary) { + t.Errorf("target after Apply = %q, want %q", got, binary) + } +} + +func TestApplyChecksumMismatch(t *testing.T) { + archive := makeTarGz(t, "pigo", []byte("real content")) + archiveName := "pigo_0.4.0_Linux_x86_64.tar.gz" + // Wrong checksum on purpose. + sums := "0000000000000000000000000000000000000000000000000000000000000000 " + archiveName + "\n" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch filepath.Base(r.URL.Path) { + case archiveName: + _, _ = w.Write(archive) + case checksumsFile: + _, _ = w.Write([]byte(sums)) + } + })) + defer srv.Close() + + dir := t.TempDir() + target := filepath.Join(dir, "pigo") + _ = os.WriteFile(target, []byte("old"), 0o755) + + u := &Updater{ + HTTPClient: srv.Client(), + ReleaseBaseURL: srv.URL, + GOOS: "linux", + GOARCH: "amd64", + ExecPath: target, + } + if err := u.Apply(context.Background(), "v0.4.0", &bytes.Buffer{}); err == nil { + t.Fatal("expected checksum mismatch error") + } + // Target must be untouched on checksum failure. + if got, _ := os.ReadFile(target); string(got) != "old" { + t.Errorf("target modified despite checksum failure: %q", got) + } +} + +func makeTarGz(t *testing.T, name string, content []byte) []byte { + t.Helper() + var buf bytes.Buffer + gz := gzip.NewWriter(&buf) + tw := tar.NewWriter(gz) + hdr := &tar.Header{Name: name, Mode: 0o755, Size: int64(len(content)), Typeflag: tar.TypeReg} + if err := tw.WriteHeader(hdr); err != nil { + t.Fatal(err) + } + if _, err := tw.Write(content); err != nil { + t.Fatal(err) + } + tw.Close() + gz.Close() + return buf.Bytes() +} diff --git a/pigo/internal/selfupdate/version.go b/pigo/internal/selfupdate/version.go new file mode 100644 index 0000000..07aa070 --- /dev/null +++ b/pigo/internal/selfupdate/version.go @@ -0,0 +1,142 @@ +// Package selfupdate provides version discovery and comparison for pigo's +// self-update feature (issue #465). It queries the GitHub Releases API for the +// latest published tag of the pigo repository and compares it against the +// build-time version injected into the main package, so both `pigo update` and +// the interactive startup banner can decide whether a newer release exists. +// +// The build version ("dev" for `go build`/`go run` from source, a real +// vX.Y.Z for goreleaser builds) is not owned by this package; callers pass it +// in. Non-release versions ("", "dev", "unknown") are treated as +// non-comparable so a source build never reports a spurious update. +package selfupdate + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "strconv" + "strings" + "time" +) + +// Repo is the GitHub "owner/name" whose releases back pigo's self-update. It +// matches the release target in .goreleaser.yaml and install.sh. +const Repo = "smallnest/pigo" + +// latestReleaseURL builds the GitHub API endpoint for a repo's latest release. +func latestReleaseURL(repo string) string { + return fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo) +} + +// release is the subset of the GitHub release JSON we consume. +type release struct { + TagName string `json:"tag_name"` +} + +// IsReleaseVersion reports whether v is a real release version that can be +// compared against a tag. The build defaults ("dev", "unknown") and the empty +// string are not release versions. +func IsReleaseVersion(v string) bool { + switch strings.TrimSpace(v) { + case "", "dev", "unknown": + return false + default: + return true + } +} + +// LatestTag queries the GitHub Releases API for repo's latest release tag +// (e.g. "v0.4.0"). If client is nil a client with a short timeout is used. When +// the GITHUB_TOKEN environment variable is set it is sent as a bearer token to +// raise the API rate limit, mirroring install.sh. +func LatestTag(ctx context.Context, client *http.Client, repo string) (string, error) { + if client == nil { + client = &http.Client{Timeout: 10 * time.Second} + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, latestReleaseURL(repo), nil) + if err != nil { + return "", fmt.Errorf("selfupdate: build request: %w", err) + } + req.Header.Set("Accept", "application/vnd.github+json") + if tok := strings.TrimSpace(os.Getenv("GITHUB_TOKEN")); tok != "" { + req.Header.Set("Authorization", "Bearer "+tok) + } + + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("selfupdate: query latest release: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("selfupdate: GitHub API returned %s", resp.Status) + } + + var rel release + if err := json.NewDecoder(resp.Body).Decode(&rel); err != nil { + return "", fmt.Errorf("selfupdate: decode release JSON: %w", err) + } + tag := strings.TrimSpace(rel.TagName) + if tag == "" { + return "", fmt.Errorf("selfupdate: latest release has empty tag_name") + } + return tag, nil +} + +// UpdateAvailable compares the current build version against a release tag and +// reports whether the tag is strictly newer. The comparable return is false +// when current is not a release version (source builds) or when either value +// cannot be parsed as a version — callers should treat non-comparable as "no +// update to offer" rather than an error. +func UpdateAvailable(current, latest string) (available, comparable bool) { + if !IsReleaseVersion(current) { + return false, false + } + cur, ok1 := parseVersion(current) + lat, ok2 := parseVersion(latest) + if !ok1 || !ok2 { + return false, false + } + return compare(lat, cur) > 0, true +} + +// parseVersion parses a semantic-ish version ("v0.4.0", "0.4.0-next") into its +// numeric major/minor/patch, ignoring any leading "v" and any pre-release or +// build suffix after "-" or "+". It reports ok=false when no numeric component +// can be read. +func parseVersion(v string) ([3]int, bool) { + s := strings.TrimSpace(v) + s = strings.TrimPrefix(s, "v") + // Drop pre-release / build metadata: "0.4.0-next" -> "0.4.0". + if i := strings.IndexAny(s, "-+"); i >= 0 { + s = s[:i] + } + if s == "" { + return [3]int{}, false + } + parts := strings.Split(s, ".") + var out [3]int + for i := 0; i < 3 && i < len(parts); i++ { + n, err := strconv.Atoi(parts[i]) + if err != nil { + return [3]int{}, false + } + out[i] = n + } + return out, true +} + +// compare returns -1, 0, or 1 as a is less than, equal to, or greater than b. +func compare(a, b [3]int) int { + for i := 0; i < 3; i++ { + switch { + case a[i] < b[i]: + return -1 + case a[i] > b[i]: + return 1 + } + } + return 0 +} diff --git a/pigo/internal/selfupdate/version_test.go b/pigo/internal/selfupdate/version_test.go new file mode 100644 index 0000000..b9b820d --- /dev/null +++ b/pigo/internal/selfupdate/version_test.go @@ -0,0 +1,142 @@ +package selfupdate + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +func TestIsReleaseVersion(t *testing.T) { + cases := map[string]bool{ + "": false, + "dev": false, + "unknown": false, + " dev ": false, + "v0.4.0": true, + "0.4.0": true, + } + for in, want := range cases { + if got := IsReleaseVersion(in); got != want { + t.Errorf("IsReleaseVersion(%q) = %v, want %v", in, got, want) + } + } +} + +func TestUpdateAvailable(t *testing.T) { + tests := []struct { + name string + current, latest string + wantAvail, wantOK bool + }{ + {"update available", "v0.3.1", "v0.4.0", true, true}, + {"patch update", "0.4.0", "0.4.1", true, true}, + {"already latest", "v0.4.0", "v0.4.0", false, true}, + {"current newer", "v0.5.0", "v0.4.0", false, true}, + {"prerelease latest", "v0.4.0", "v0.4.1-next", true, true}, + {"dev current not comparable", "dev", "v0.4.0", false, false}, + {"unknown current not comparable", "unknown", "v0.4.0", false, false}, + {"unparseable latest", "v0.4.0", "not-a-version", false, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + avail, ok := UpdateAvailable(tt.current, tt.latest) + if avail != tt.wantAvail || ok != tt.wantOK { + t.Errorf("UpdateAvailable(%q,%q) = (%v,%v), want (%v,%v)", + tt.current, tt.latest, avail, ok, tt.wantAvail, tt.wantOK) + } + }) + } +} + +func TestParseVersion(t *testing.T) { + tests := []struct { + in string + want [3]int + ok bool + }{ + {"v1.2.3", [3]int{1, 2, 3}, true}, + {"1.2.3", [3]int{1, 2, 3}, true}, + {"0.4.0-next", [3]int{0, 4, 0}, true}, + {"1.2", [3]int{1, 2, 0}, true}, + {"", [3]int{}, false}, + {"vabc", [3]int{}, false}, + } + for _, tt := range tests { + got, ok := parseVersion(tt.in) + if got != tt.want || ok != tt.ok { + t.Errorf("parseVersion(%q) = (%v,%v), want (%v,%v)", tt.in, got, ok, tt.want, tt.ok) + } + } +} + +func TestLatestTag(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Accept") != "application/vnd.github+json" { + t.Errorf("missing Accept header") + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"tag_name":"v0.4.0","name":"pigo 0.4.0"}`)) + })) + defer srv.Close() + + // LatestTag builds the URL from repo; use a transport that redirects to the + // test server regardless of host. + client := srv.Client() + client.Transport = rewriteHost{base: srv.URL, rt: client.Transport} + + tag, err := LatestTag(context.Background(), client, "smallnest/pigo") + if err != nil { + t.Fatalf("LatestTag: %v", err) + } + if tag != "v0.4.0" { + t.Errorf("tag = %q, want v0.4.0", tag) + } +} + +func TestLatestTagErrors(t *testing.T) { + t.Run("non-200", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + defer srv.Close() + client := srv.Client() + client.Transport = rewriteHost{base: srv.URL, rt: client.Transport} + if _, err := LatestTag(context.Background(), client, "smallnest/pigo"); err == nil { + t.Fatal("expected error on 403") + } + }) + + t.Run("empty tag", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"tag_name":""}`)) + })) + defer srv.Close() + client := srv.Client() + client.Transport = rewriteHost{base: srv.URL, rt: client.Transport} + if _, err := LatestTag(context.Background(), client, "smallnest/pigo"); err == nil { + t.Fatal("expected error on empty tag_name") + } + }) +} + +// rewriteHost redirects every request to base, so tests can point the fixed +// GitHub API URL at an httptest server. +type rewriteHost struct { + base string + rt http.RoundTripper +} + +func (rw rewriteHost) RoundTrip(req *http.Request) (*http.Response, error) { + u, err := req.URL.Parse(rw.base) + if err != nil { + return nil, err + } + req.URL.Scheme = u.Scheme + req.URL.Host = u.Host + rt := rw.rt + if rt == nil { + rt = http.DefaultTransport + } + return rt.RoundTrip(req) +} diff --git a/pigo/internal/session/export.go b/pigo/internal/session/export.go new file mode 100644 index 0000000..d1a1e22 --- /dev/null +++ b/pigo/internal/session/export.go @@ -0,0 +1,112 @@ +// This file implements session export/import (US-008, #124): a session can be +// exported to a self-contained JSONL or HTML file, and a JSONL export can be +// imported back as a fresh, resumable session. The JSONL form is the same +// role-discriminated schema the store persists, so an export → import round-trip +// is lossless (message contents and the id/parentId tree survive verbatim); the +// HTML form is a read-only, self-contained transcript with inline styles and no +// external network resources, suitable for sharing. +package session + +import ( + "fmt" + "io" + "os" + "path/filepath" + "strings" + "time" +) + +// WriteJSONL writes header + entries as JSONL in the store's on-disk schema +// (header line first, then one entry line each, ids/parentIds preserved). It is +// the export counterpart to writeSessionEntries and the exact input ReadJSONL +// expects, so a WriteJSONL → ReadJSONL round-trip is lossless. +func WriteJSONL(w io.Writer, header SessionHeader, entries []Entry) error { + header.Version = SchemaVersion + return writeSessionEntries(w, header, entries) +} + +// ReadJSONL decodes a JSONL export (as produced by WriteJSONL or a raw session +// file) into a header and entries, migrating v1/v2 bare-message files the same +// way LoadEntries does. It is the import counterpart to WriteJSONL. +func ReadJSONL(r io.Reader) (SessionHeader, []Entry, error) { + return readSession(r) +} + +// Export writes the session identified by id to outPath. The format is chosen +// by outPath's extension: ".html"/".htm" produces a self-contained HTML +// transcript; anything else (including ".jsonl") produces JSONL. The parent +// directory of outPath must already exist. It returns the number of entries +// written so a caller can report progress. +func (s *Store) Export(id, outPath string) (int, error) { + header, entries, err := s.LoadEntries(id) + if err != nil { + return 0, err + } + f, err := os.Create(outPath) + if err != nil { + return 0, fmt.Errorf("session: create export %s: %w", outPath, err) + } + defer f.Close() + ext := strings.ToLower(filepath.Ext(outPath)) + if ext == ".html" || ext == ".htm" { + if err := WriteHTML(f, header, entries); err != nil { + return 0, err + } + } else { + if err := WriteJSONL(f, header, entries); err != nil { + return 0, err + } + } + if err := f.Close(); err != nil { + return 0, fmt.Errorf("session: finalize export %s: %w", outPath, err) + } + return len(entries), nil +} + +// Import reads a JSONL export at inPath and materializes it as a fresh session +// in the store: a new id (derived from now) is assigned, the original id is +// recorded as ParentSession for lineage, and the entries are written verbatim +// (ids/parentIds preserved) so the tree — and thus PathToLeaf/resume — behaves +// exactly as in the source. It returns the new header and the imported entries. +// An HTML file (or any non-JSONL input) fails to parse and returns an error +// rather than importing garbage. +func (s *Store) Import(inPath string, now time.Time) (SessionHeader, []Entry, error) { + f, err := os.Open(inPath) + if err != nil { + return SessionHeader{}, nil, fmt.Errorf("session: open import %s: %w", inPath, err) + } + defer f.Close() + srcHeader, entries, err := ReadJSONL(f) + if err != nil { + return SessionHeader{}, nil, err + } + newHeader := SessionHeader{ + ID: NewID(now), + CreatedAt: now, + UpdatedAt: now, + Model: srcHeader.Model, + Provider: srcHeader.Provider, + SystemPrompt: srcHeader.SystemPrompt, + ParentSession: srcHeader.ID, + } + if err := s.SaveEntries(newHeader, entries); err != nil { + return SessionHeader{}, nil, err + } + return newHeader, entries, nil +} + +// WriteHTML writes a self-contained HTML transcript of the session: inline CSS +// only (no external stylesheets, fonts, scripts, or network resources), role +// color-coding, and tool-call/result blocks. All message text is HTML-escaped +// so a transcript containing markup or a crafted "" cannot break out of +// its container or inject active content (defensive against a hostile session). +func WriteHTML(w io.Writer, header SessionHeader, entries []Entry) error { + var b strings.Builder + b.WriteString(htmlHead(header)) + for _, e := range entries { + b.WriteString(renderEntryHTML(e)) + } + b.WriteString(htmlFoot()) + _, err := io.WriteString(w, b.String()) + return err +} diff --git a/pigo/internal/session/export_html.go b/pigo/internal/session/export_html.go new file mode 100644 index 0000000..33cd1eb --- /dev/null +++ b/pigo/internal/session/export_html.go @@ -0,0 +1,125 @@ +// This file holds the HTML rendering helpers for session export (US-008, #124). +// The output is a single self-contained document: all CSS is inlined in a +// + + +
+

pigo session — %s

+

%s

+`, html.EscapeString(title), html.EscapeString(title), strings.Join(meta, " · ")) +} + +// htmlFoot closes the container and document. +func htmlFoot() string { + return ` +
+ + +` +} + +// renderEntryHTML renders one entry as a role-colored message block. Text and +// tool arguments are escaped so no session content can inject markup. +func renderEntryHTML(e Entry) string { + switch m := e.Message.(type) { + case agentcore.UserMessage: + return msgBlock("user", "User", html.EscapeString(agentcore.ContentToText(m.Content)), "") + case agentcore.AssistantMessage: + var tools strings.Builder + for _, c := range m.ToolCalls() { + args := strings.TrimSpace(string(c.Arguments)) + tools.WriteString(fmt.Sprintf(`
→ %s %s
`, + html.EscapeString(c.Name), html.EscapeString(args))) + } + return msgBlock("assistant", "Assistant", html.EscapeString(agentcore.ContentToText(m.Content)), tools.String()) + case agentcore.ToolResultMessage: + label := "Tool Result" + if m.ToolName != "" { + label = "Tool Result: " + m.ToolName + } + return msgBlock("tool", html.EscapeString(label), html.EscapeString(agentcore.ContentToText(m.Content)), "") + case agentcore.CompactionMessage: + return msgBlock("compaction", "Compaction", html.EscapeString(m.Summary), "") + default: + return msgBlock("assistant", html.EscapeString(e.Message.Role()), "", "") + } +} + +// msgBlock assembles one .msg block with a role class, a role label, the escaped +// body text, and optional pre-rendered tool-call HTML. Callers MUST pass already +// escaped text/label; extra is trusted HTML built here from escaped parts. +func msgBlock(class, label, escapedText, extra string) string { + var b strings.Builder + b.WriteString(`
`) + b.WriteString(label) + b.WriteString(`
`) + if escapedText != "" { + b.WriteString(`
`) + b.WriteString(escapedText) + b.WriteString(`
`) + } + b.WriteString(extra) + b.WriteString("
\n") + return b.String() +} diff --git a/pigo/internal/session/export_test.go b/pigo/internal/session/export_test.go new file mode 100644 index 0000000..0328218 --- /dev/null +++ b/pigo/internal/session/export_test.go @@ -0,0 +1,197 @@ +package session + +// Tests for session export/import (US-008, #124). They cover the lossless +// JSONL round-trip (export → import preserves the header fields, message +// sequence, and entry tree), the self-contained HTML export (inline styles, no +// external network resources, HTML-escaped content), and format selection by +// file extension. + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/smallnest/pigo/internal/agentcore" +) + +// seedSession writes a small multi-turn session and returns its header. +func seedSession(t *testing.T, s *Store) SessionHeader { + t.Helper() + now := time.Date(2026, 7, 17, 9, 0, 0, 0, time.UTC) + header := SessionHeader{ + ID: NewID(now), + CreatedAt: now, + UpdatedAt: now, + Model: "anthropic/claude-opus-4", + Provider: "anthropic", + SystemPrompt: "You are pigo.", + } + if err := s.Save(header, sampleMessages()); err != nil { + t.Fatalf("Save: %v", err) + } + return header +} + +// TestExportImportJSONLRoundTrip is the core acceptance check: exporting a +// session to JSONL and importing it back yields the same message sequence and +// header carry-over, with the source id recorded as the import's parent. +func TestExportImportJSONLRoundTrip(t *testing.T) { + s := newStore(t) + header := seedSession(t, s) + + out := filepath.Join(t.TempDir(), "export.jsonl") + n, err := s.Export(header.ID, out) + if err != nil { + t.Fatalf("Export: %v", err) + } + if n != len(sampleMessages()) { + t.Errorf("exported %d entries, want %d", n, len(sampleMessages())) + } + + now := time.Date(2026, 7, 17, 10, 0, 0, 0, time.UTC) + newHeader, entries, err := s.Import(out, now) + if err != nil { + t.Fatalf("Import: %v", err) + } + if newHeader.ID == header.ID { + t.Errorf("import must assign a fresh id, got the source id %q", newHeader.ID) + } + if newHeader.ParentSession != header.ID { + t.Errorf("ParentSession = %q, want source id %q", newHeader.ParentSession, header.ID) + } + if newHeader.Model != header.Model || newHeader.Provider != header.Provider || newHeader.SystemPrompt != header.SystemPrompt { + t.Errorf("header fields not carried over: %+v", newHeader) + } + if len(entries) != len(sampleMessages()) { + t.Fatalf("imported %d entries, want %d", len(entries), len(sampleMessages())) + } + + // The imported session must be independently loadable with the same roles. + _, gotMsgs, err := s.Load(newHeader.ID) + if err != nil { + t.Fatalf("Load imported: %v", err) + } + wantRoles := []string{agentcore.RoleUser, agentcore.RoleAssistant, agentcore.RoleToolResult, agentcore.RoleAssistant} + if len(gotMsgs) != len(wantRoles) { + t.Fatalf("imported message count = %d, want %d", len(gotMsgs), len(wantRoles)) + } + for i, m := range gotMsgs { + if m.Role() != wantRoles[i] { + t.Errorf("message[%d] role = %q, want %q", i, m.Role(), wantRoles[i]) + } + } + // The tool call must survive the round-trip. + a, ok := gotMsgs[1].(agentcore.AssistantMessage) + if !ok { + t.Fatalf("message[1] is not AssistantMessage: %T", gotMsgs[1]) + } + if calls := a.ToolCalls(); len(calls) != 1 || calls[0].Name != "read" { + t.Errorf("tool calls = %+v, want one 'read'", calls) + } +} + +// TestWriteReadJSONLPreservesTree checks the lower-level primitives: WriteJSONL +// then ReadJSONL preserves entry ids and parentIds verbatim (so the tree — and +// thus resume/PathToLeaf — behaves identically). +func TestWriteReadJSONLPreservesTree(t *testing.T) { + s := newStore(t) + header := seedSession(t, s) + _, entries, err := s.LoadEntries(header.ID) + if err != nil { + t.Fatalf("LoadEntries: %v", err) + } + + var buf bytes.Buffer + if err := WriteJSONL(&buf, header, entries); err != nil { + t.Fatalf("WriteJSONL: %v", err) + } + gotHeader, gotEntries, err := ReadJSONL(&buf) + if err != nil { + t.Fatalf("ReadJSONL: %v", err) + } + if gotHeader.Version != SchemaVersion { + t.Errorf("version = %d, want %d", gotHeader.Version, SchemaVersion) + } + if len(gotEntries) != len(entries) { + t.Fatalf("entry count = %d, want %d", len(gotEntries), len(entries)) + } + for i := range entries { + if gotEntries[i].ID != entries[i].ID || gotEntries[i].ParentID != entries[i].ParentID { + t.Errorf("entry[%d] tree ids changed: got {%q,%q} want {%q,%q}", + i, gotEntries[i].ID, gotEntries[i].ParentID, entries[i].ID, entries[i].ParentID) + } + } +} + +// TestExportHTMLSelfContained verifies the HTML export is a self-contained +// document: it carries inline styles, has no external network resources (no +// http(s):// URLs, no ")}}, + agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, Content: agentcore.ContentList{agentcore.NewTextContent("safe & sound")}, StopReason: agentcore.StopReasonEndTurn}, + } + if err := s.Save(header, msgs); err != nil { + t.Fatalf("Save: %v", err) + } + + out := filepath.Join(t.TempDir(), "export.html") + if _, err := s.Export(header.ID, out); err != nil { + t.Fatalf("Export html: %v", err) + } + data, err := os.ReadFile(out) + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + doc := string(data) + + if !strings.Contains(doc, "