diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index 2d544d4f0..795fa5eba 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -41,10 +41,11 @@ jobs: with: go-version-file: go.mod + - name: Install govulncheck + run: go install golang.org/x/vuln/cmd/govulncheck@v1.1.4 + - name: Run Govulncheck - uses: golang/govulncheck-action@v1 - with: - go-package: ./... + run: govulncheck -C . -format text ./... test: name: Tests diff --git a/assets/wechat.png b/assets/wechat.png index 07a05dd91..66ffa99e9 100644 Binary files a/assets/wechat.png and b/assets/wechat.png differ diff --git a/docker/Dockerfile.goreleaser.launcher b/docker/Dockerfile.goreleaser.launcher index 5d65576f7..0a20a90b3 100644 --- a/docker/Dockerfile.goreleaser.launcher +++ b/docker/Dockerfile.goreleaser.launcher @@ -9,4 +9,4 @@ COPY $TARGETPLATFORM/picoclaw-launcher /usr/local/bin/picoclaw-launcher COPY $TARGETPLATFORM/picoclaw-launcher-tui /usr/local/bin/picoclaw-launcher-tui ENTRYPOINT ["picoclaw-launcher"] -CMD ["-public", "-no-browser"] +CMD ["-console", "-public", "-no-browser"] diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 0bf46a2ae..7c940621f 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -45,8 +45,11 @@ services: - launcher environment: - PICOCLAW_GATEWAY_HOST=0.0.0.0 + # Set a fixed dashboard token instead of a random one each restart. + # If not set, a random token is generated and printed to the console on startup. + #- PICOCLAW_LAUNCHER_TOKEN=your-secret-token-here ports: - - "127.0.0.1:18800:18800" - - "127.0.0.1:18790:18790" + - "18800:18800" + - "18790:18790" volumes: - ./data:/root/.picoclaw diff --git a/docs/hooks/README.md b/docs/hooks/README.md index ec3bbc46a..5be0f30b5 100644 --- a/docs/hooks/README.md +++ b/docs/hooks/README.md @@ -28,6 +28,69 @@ The currently exposed synchronous hook points are: Everything else is exposed as read-only events. +## Hook Actions + +Hooks can return different actions to control the flow: + +| Action | Applicable Stages | Effect | +| --- | --- | --- | +| `continue` | All interceptors | Pass through without modification | +| `modify` | `before_llm`, `after_llm`, `before_tool`, `after_tool` | Modify request/response and continue | +| `respond` | `before_tool` | Return a tool result directly, skip actual tool execution | +| `deny_tool` | `before_tool` | Deny tool execution, return error message | +| `abort_turn` | All interceptors | Abort the current turn | +| `hard_abort` | All interceptors | Force stop the entire agent loop | + +### The `respond` Action + +The `respond` action is special: it allows a `before_tool` hook to provide the tool result directly, skipping the actual tool execution. This is useful for: + +1. **Plugin tool injection**: External hooks can implement tools without registering them in the tool registry +2. **Tool result caching**: Return cached results for repeated tool calls +3. **Tool mocking**: Return mock results for testing purposes + +When a hook returns `respond` with a `HookResult`, the agent loop: +1. Skips the actual tool execution +2. Uses the provided result as if the tool had executed +3. Continues the turn normally with the result + +Example (Go in-process hook): + +```go +func (h *MyHook) BeforeTool( + ctx context.Context, + call *agent.ToolCallHookRequest, +) (*agent.ToolCallHookRequest, agent.HookDecision, error) { + if call.Tool == "my_plugin_tool" { + next := call.Clone() + next.HookResult = &tools.ToolResult{ + ForLLM: "Plugin tool executed successfully", + Silent: false, + IsError: false, + } + return next, agent.HookDecision{Action: agent.HookActionRespond}, nil + } + return call, agent.HookDecision{Action: agent.HookActionContinue}, nil +} +``` + +Example (Python process hook): + +```python +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + if tool == "my_plugin_tool": + return { + "action": "respond", + "result": { + "for_llm": "Plugin tool executed successfully", + "silent": False, + "is_error": False + } + } + return {"action": "continue"} +``` + ## Execution Order `HookManager` sorts hooks like this: diff --git a/docs/hooks/README.zh.md b/docs/hooks/README.zh.md index 46c7c9392..2170d45c8 100644 --- a/docs/hooks/README.zh.md +++ b/docs/hooks/README.zh.md @@ -28,6 +28,69 @@ 其余 lifecycle 通过事件形式只读暴露。 +## Hook Actions + +Hook 可以返回不同的 action 来控制流程： + +| Action | 适用阶段 | 效果 | +| --- | --- | --- | +| `continue` | 所有拦截型 | 放行，不做修改 | +| `modify` | `before_llm`, `after_llm`, `before_tool`, `after_tool` | 改写请求/响应后放行 | +| `respond` | `before_tool` | 直接返回工具结果，跳过实际工具执行 | +| `deny_tool` | `before_tool` | 拒绝工具执行，返回错误信息 | +| `abort_turn` | 所有拦截型 | 中止当前 turn | +| `hard_abort` | 所有拦截型 | 强制终止整个 agent loop | + +### `respond` Action + +`respond` action 是特殊的：它允许 `before_tool` hook 直接提供工具结果，跳过实际工具执行。适用于： + +1. **插件工具注入**：外部 hook 可以实现工具，无需在 ToolRegistry 注册 +2. **工具结果缓存**：对重复调用返回缓存结果 +3. **工具模拟**：测试时返回模拟结果 + +当 hook 返回 `respond` 并携带 `HookResult` 时，agent loop 会： +1. 跳过实际工具执行 +2. 使用提供的结果作为工具执行结果 +3. 正常继续 turn 流程 + +示例（Go 进程内 hook）： + +```go +func (h *MyHook) BeforeTool( + ctx context.Context, + call *agent.ToolCallHookRequest, +) (*agent.ToolCallHookRequest, agent.HookDecision, error) { + if call.Tool == "my_plugin_tool" { + next := call.Clone() + next.HookResult = &tools.ToolResult{ + ForLLM: "Plugin tool executed successfully", + Silent: false, + IsError: false, + } + return next, agent.HookDecision{Action: agent.HookActionRespond}, nil + } + return call, agent.HookDecision{Action: agent.HookActionContinue}, nil +} +``` + +示例（Python process hook）： + +```python +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + if tool == "my_plugin_tool": + return { + "action": "respond", + "result": { + "for_llm": "Plugin tool executed successfully", + "silent": False, + "is_error": False + } + } + return {"action": "continue"} +``` + ## 执行顺序 HookManager 的排序规则是： diff --git a/docs/hooks/hook-json-protocol.md b/docs/hooks/hook-json-protocol.md new file mode 100644 index 000000000..58b6e323b --- /dev/null +++ b/docs/hooks/hook-json-protocol.md @@ -0,0 +1,568 @@ +# Hook JSON-RPC Protocol Details + +All hooks use `JSON-RPC 2.0` format, with one JSON message per line, transmitted via stdio. + +--- + +## Basic Protocol Structure + +### Request (PicoClaw → Hook) + +```json +{"jsonrpc":"2.0","id":1,"method":"hook.xxx","params":{...}} +``` + +### Response (Hook → PicoClaw) + +Success: +```json +{"jsonrpc":"2.0","id":1,"result":{...}} +``` + +Error: +```json +{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"error message"}} +``` + +--- + +## 1. `hook.hello` (Handshake) + +Handshake must be completed at startup, otherwise the hook process will be terminated. + +### Request + +```json +{ + "jsonrpc": "2.0", + "id": 1, + "method": "hook.hello", + "params": { + "name": "py_review_gate", + "version": 1, + "modes": ["observe", "tool", "approve"] + } +} +``` + +| Field | Description | +|-------|-------------| +| `name` | hook name (from configuration) | +| `version` | protocol version, currently `1` | +| `modes` | capability modes supported by the hook | + +### Response + +```json +{ + "jsonrpc": "2.0", + "id": 1, + "result": { + "ok": true, + "name": "python-review-gate" + } +} +``` + +--- + +## 2. `hook.before_llm` + +Triggered before sending request to LLM. Can be used to inject tools. + +### Request + +```json +{ + "jsonrpc": "2.0", + "id": 2, + "method": "hook.before_llm", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "ParentTurnID": "", + "SessionKey": "session-1", + "Iteration": 0, + "TracePath": "runTurn", + "Source": "turn.llm.request" + }, + "model": "claude-sonnet", + "messages": [ + {"role": "user", "content": "hello"} + ], + "tools": [ + { + "type": "function", + "function": { + "name": "echo", + "description": "echo text", + "parameters": {"type": "object"} + } + } + ], + "options": { + "temperature": 0.7 + }, + "channel": "cli", + "chat_id": "chat-1", + "graceful_terminal": false + } +} +``` + +| Field | Description | +|-------|-------------| +| `meta` | event metadata for tracing | +| `model` | requested model name | +| `messages` | conversation history | +| `tools` | list of available tool definitions | +| `options` | LLM parameters (temperature, max_tokens, etc.) | +| `channel` | request source channel | +| `chat_id` | session ID | + +### Response (Tool Injection Example) + +```json +{ + "jsonrpc": "2.0", + "id": 2, + "result": { + "action": "modify", + "request": { + "model": "claude-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "function": { + "name": "echo", + "description": "echo", + "parameters": {} + } + }, + { + "type": "function", + "function": { + "name": "my_plugin_tool", + "description": "Plugin injected tool", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string"} + } + } + } + } + ] + } + } +} +``` + +| Field | Description | +|-------|-------------| +| `action` | decision action (see table below) | +| `request` | modified request object | + +--- + +## 3. `hook.after_llm` + +Triggered after receiving LLM response. Can modify response content. + +### Request + +```json +{ + "jsonrpc": "2.0", + "id": 3, + "method": "hook.after_llm", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "model": "claude-sonnet", + "response": { + "role": "assistant", + "content": "Hi!", + "tool_calls": [ + { + "id": "tc-1", + "type": "function", + "function": { + "name": "echo", + "arguments": "{\"text\":\"hi\"}" + } + } + ] + }, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +### Response + +```json +{ + "jsonrpc": "2.0", + "id": 3, + "result": { + "action": "continue" + } +} +``` + +--- + +## 4. `hook.before_tool` + +Triggered before tool execution. Can modify tool name and arguments, deny execution, or return result directly. + +### Request + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "method": "hook.before_tool", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "tool": "echo_text", + "arguments": { + "text": "hello" + }, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +| Field | Description | +|-------|-------------| +| `tool` | tool name | +| `arguments` | tool arguments | + +### Response (Modify Arguments) + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "action": "modify", + "call": { + "tool": "echo_text", + "arguments": { + "text": "modified hello" + } + } + } +} +``` + +### Response (Deny Execution) + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "action": "deny_tool", + "reason": "Invalid arguments" + } +} +``` + +### Response (Return Result Directly - respond) + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "action": "respond", + "call": { + "tool": "my_plugin_tool", + "arguments": { + "query": "hello" + } + }, + "result": { + "for_llm": "Plugin tool executed successfully", + "for_user": "", + "silent": false, + "is_error": false + } + } +} +``` + +The `respond` action allows hooks to return tool results directly, skipping actual tool execution. Use cases: +1. **Plugin tool injection**: External hooks can implement tools without registering in ToolRegistry +2. **Tool result caching**: Return cached results for repeated calls +3. **Tool mocking**: Return mock results during testing + +| Field | Description | +|-------|-------------| +| `action` | must be `respond` | +| `call` | modified call information (optional) | +| `result` | tool result to return directly | + +--- + +## 5. `hook.after_tool` + +Triggered after tool execution completes. Can modify the result returned to LLM. + +### Request + +```json +{ + "jsonrpc": "2.0", + "id": 5, + "method": "hook.after_tool", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "tool": "echo_text", + "arguments": { + "text": "hello" + }, + "result": { + "for_llm": "echoed: hello", + "for_user": "", + "silent": false, + "is_error": false, + "async": false, + "media": [], + "artifact_tags": [], + "response_handled": false + }, + "duration": 15000000, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +| Field | Description | +|-------|-------------| +| `result.for_llm` | content returned to LLM | +| `result.for_user` | content sent to user | +| `result.silent` | whether silent (not sent to user) | +| `result.is_error` | whether it's an error | +| `result.async` | whether executed asynchronously | +| `result.media` | list of media references | +| `result.artifact_tags` | local artifact path tags | +| `result.response_handled` | whether response has been handled | +| `duration` | execution time (nanoseconds) | + +### Response + +```json +{ + "jsonrpc": "2.0", + "id": 5, + "result": { + "action": "continue" + } +} +``` + +--- + +## 6. `hook.approve_tool` + +Approval hook for deciding whether to allow execution of sensitive tools. + +### Request + +```json +{ + "jsonrpc": "2.0", + "id": 6, + "method": "hook.approve_tool", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "tool": "bash", + "arguments": { + "command": "rm -rf /" + }, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +### Response (Approved) + +```json +{ + "jsonrpc": "2.0", + "id": 6, + "result": { + "approved": true + } +} +``` + +### Response (Denied) + +```json +{ + "jsonrpc": "2.0", + "id": 6, + "result": { + "approved": false, + "reason": "Dangerous command, execution denied" + } +} +``` + +--- + +## 7. `hook.event` (notification) + +Observer event, broadcast only, no response required. `id` is `0` or absent. + +```json +{ + "jsonrpc": "2.0", + "method": "hook.event", + "params": { + "Kind": "tool_exec_start", + "Meta": { + "AgentID": "agent-1", + "TurnID": "turn-1" + }, + "Payload": { + "Tool": "echo_text", + "Arguments": {"text": "hello"} + } + } +} +``` + +Common `Kind` values: +- `turn_start` / `turn_end` +- `llm_request` / `llm_response` +- `tool_exec_start` / `tool_exec_end` / `tool_exec_skipped` +- `steering_injected` +- `interrupt_received` +- `error` + +--- + +## Action Options + +| action | Applicable hooks | Effect | +|--------|-----------------|--------| +| `continue` | All interceptor types | Pass through without modification | +| `modify` | `before_llm`, `before_tool`, `after_llm`, `after_tool` | Modify request/response and pass through | +| `respond` | `before_tool` | Return tool result directly, skip actual execution. **Note: AfterTool is NOT called (design decision - respond provides final answer).** | +| `deny_tool` | `before_tool` | Deny tool execution | +| `abort_turn` | All interceptor types | Abort current turn, return error | +| `hard_abort` | All interceptor types | Force stop entire agent loop | + +--- + +## Complete Flow Example + +```json +{"jsonrpc":"2.0","id":1,"method":"hook.hello","params":{"name":"my_hook","version":1,"modes":["tool","approve"]}} +{"jsonrpc":"2.0","id":1,"result":{"ok":true,"name":"my_hook"}} +{"jsonrpc":"2.0","id":2,"method":"hook.before_llm","params":{"model":"claude-sonnet","messages":[{"role":"user","content":"hello"}],"tools":[]}} +{"jsonrpc":"2.0","id":2,"result":{"action":"continue"}} +{"jsonrpc":"2.0","id":3,"method":"hook.before_tool","params":{"tool":"bash","arguments":{"command":"ls"}}} +{"jsonrpc":"2.0","id":3,"result":{"action":"continue"}} +{"jsonrpc":"2.0","id":4,"method":"hook.approve_tool","params":{"tool":"bash","arguments":{"command":"ls"}}} +{"jsonrpc":"2.0","id":4,"result":{"approved":true}} +{"jsonrpc":"2.0","id":5,"method":"hook.after_tool","params":{"tool":"bash","arguments":{"command":"ls"},"result":{"for_llm":"file1.txt\nfile2.txt"},"duration":5000000}} +{"jsonrpc":"2.0","id":5,"result":{"action":"continue"}} +{"jsonrpc":"2.0","id":6,"method":"hook.after_llm","params":{"model":"claude-sonnet","response":{"role":"assistant","content":"Files listed"}}} +{"jsonrpc":"2.0","id":6,"result":{"action":"continue"}} +``` + +--- + +## Plugin Tool Injection via `before_llm` and `before_tool` + +Standard flow for plugin tool injection: + +1. In `before_llm`, inject tool definition to let LLM know the tool is available +2. In `before_tool`, use `respond` action to return tool execution result directly + +### `before_llm` Inject Tool Definition + +```python +def handle_before_llm(params: dict) -> dict: + tools = params.get("tools", []) + + # Add plugin tool definition + tools.append({ + "type": "function", + "function": { + "name": "my_plugin_tool", + "description": "Plugin provided tool", + "parameters": { + "type": "object", + "properties": { + "input": {"type": "string", "description": "Input content"} + }, + "required": ["input"] + } + } + }) + + return { + "action": "modify", + "request": { + "model": params["model"], + "messages": params["messages"], + "tools": tools, + "options": params.get("options", {}) + } + } +``` + +### `before_tool` Return Execution Result + +```python +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + + if tool == "my_plugin_tool": + # Implement tool logic here + args = params.get("arguments", {}) + input_text = args.get("input", "") + + # Return result directly, no need to register in ToolRegistry + return { + "action": "respond", + "result": { + "for_llm": f"Plugin tool executed successfully, input: {input_text}", + "silent": False, + "is_error": False + } + } + + return {"action": "continue"} +``` + +This way, external hooks can fully implement plugin tools without registering any tool implementation inside PicoClaw. \ No newline at end of file diff --git a/docs/hooks/hook-json-protocol.zh.md b/docs/hooks/hook-json-protocol.zh.md new file mode 100644 index 000000000..675e0a429 --- /dev/null +++ b/docs/hooks/hook-json-protocol.zh.md @@ -0,0 +1,568 @@ +# Hook JSON-RPC 协议详解 + +所有 hook 使用 `JSON-RPC 2.0` 格式，每行一个 JSON 消息，通过 stdio 传输。 + +--- + +## 基础协议结构 + +### 请求（PicoClaw → Hook） + +```json +{"jsonrpc":"2.0","id":1,"method":"hook.xxx","params":{...}} +``` + +### 响应（Hook → PicoClaw） + +成功： +```json +{"jsonrpc":"2.0","id":1,"result":{...}} +``` + +错误： +```json +{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"错误信息"}} +``` + +--- + +## 1. `hook.hello`（握手） + +启动时必须完成握手，否则 hook 进程会被终止。 + +### 请求 + +```json +{ + "jsonrpc": "2.0", + "id": 1, + "method": "hook.hello", + "params": { + "name": "py_review_gate", + "version": 1, + "modes": ["observe", "tool", "approve"] + } +} +``` + +| 字段 | 说明 | +|------|------| +| `name` | hook 名称（来自配置） | +| `version` | 协议版本，当前为 `1` | +| `modes` | hook 支持的能力模式 | + +### 响应 + +```json +{ + "jsonrpc": "2.0", + "id": 1, + "result": { + "ok": true, + "name": "python-review-gate" + } +} +``` + +--- + +## 2. `hook.before_llm` + +在发送请求给 LLM 之前触发。可用于注入工具。 + +### 请求 + +```json +{ + "jsonrpc": "2.0", + "id": 2, + "method": "hook.before_llm", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "ParentTurnID": "", + "SessionKey": "session-1", + "Iteration": 0, + "TracePath": "runTurn", + "Source": "turn.llm.request" + }, + "model": "claude-sonnet", + "messages": [ + {"role": "user", "content": "hello"} + ], + "tools": [ + { + "type": "function", + "function": { + "name": "echo", + "description": "echo text", + "parameters": {"type": "object"} + } + } + ], + "options": { + "temperature": 0.7 + }, + "channel": "cli", + "chat_id": "chat-1", + "graceful_terminal": false + } +} +``` + +| 字段 | 说明 | +|------|------| +| `meta` | 事件元数据，用于追踪 | +| `model` | 请求的模型名称 | +| `messages` | 对话历史 | +| `tools` | 可用工具定义列表 | +| `options` | LLM 参数（temperature、max_tokens 等） | +| `channel` | 请求来源通道 | +| `chat_id` | 会话 ID | + +### 响应（注入工具示例） + +```json +{ + "jsonrpc": "2.0", + "id": 2, + "result": { + "action": "modify", + "request": { + "model": "claude-sonnet", + "messages": [{"role": "user", "content": "hello"}], + "tools": [ + { + "type": "function", + "function": { + "name": "echo", + "description": "echo", + "parameters": {} + } + }, + { + "type": "function", + "function": { + "name": "my_plugin_tool", + "description": "插件注入的工具", + "parameters": { + "type": "object", + "properties": { + "query": {"type": "string"} + } + } + } + } + ] + } + } +} +``` + +| 字段 | 说明 | +|------|------| +| `action` | 决策动作（见下表） | +| `request` | 修改后的请求对象 | + +--- + +## 3. `hook.after_llm` + +在收到 LLM 响应后触发。可修改响应内容。 + +### 请求 + +```json +{ + "jsonrpc": "2.0", + "id": 3, + "method": "hook.after_llm", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "model": "claude-sonnet", + "response": { + "role": "assistant", + "content": "Hi!", + "tool_calls": [ + { + "id": "tc-1", + "type": "function", + "function": { + "name": "echo", + "arguments": "{\"text\":\"hi\"}" + } + } + ] + }, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +### 响应 + +```json +{ + "jsonrpc": "2.0", + "id": 3, + "result": { + "action": "continue" + } +} +``` + +--- + +## 4. `hook.before_tool` + +在执行工具前触发。可修改工具名称和参数，或拒绝执行，或直接返回结果。 + +### 请求 + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "method": "hook.before_tool", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "tool": "echo_text", + "arguments": { + "text": "hello" + }, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +| 字段 | 说明 | +|------|------| +| `tool` | 工具名称 | +| `arguments` | 工具参数 | + +### 响应（改写参数） + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "action": "modify", + "call": { + "tool": "echo_text", + "arguments": { + "text": "modified hello" + } + } + } +} +``` + +### 响应（拒绝执行） + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "action": "deny_tool", + "reason": "参数不合法" + } +} +``` + +### 响应（直接返回结果 - respond） + +```json +{ + "jsonrpc": "2.0", + "id": 4, + "result": { + "action": "respond", + "call": { + "tool": "my_plugin_tool", + "arguments": { + "query": "hello" + } + }, + "result": { + "for_llm": "Plugin tool executed successfully", + "for_user": "", + "silent": false, + "is_error": false + } + } +} +``` + +`respond` action 允许 hook 直接返回工具结果，跳过实际工具执行。适用于： +1. **插件工具注入**：外部 hook 可实现工具，无需在 ToolRegistry 注册 +2. **工具结果缓存**：对重复调用返回缓存结果 +3. **工具模拟**：测试时返回模拟结果 + +| 字段 | 说明 | +|------|------| +| `action` | 必须为 `respond` | +| `call` | 修改后的调用信息（可选） | +| `result` | 直接返回的工具结果 | + +--- + +## 5. `hook.after_tool` + +在工具执行完成后触发。可修改返回给 LLM 的结果。 + +### 请求 + +```json +{ + "jsonrpc": "2.0", + "id": 5, + "method": "hook.after_tool", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "tool": "echo_text", + "arguments": { + "text": "hello" + }, + "result": { + "for_llm": "echoed: hello", + "for_user": "", + "silent": false, + "is_error": false, + "async": false, + "media": [], + "artifact_tags": [], + "response_handled": false + }, + "duration": 15000000, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +| 字段 | 说明 | +|------|------| +| `result.for_llm` | 返回给 LLM 的内容 | +| `result.for_user` | 发送给用户的内容 | +| `result.silent` | 是否静默（不发送给用户） | +| `result.is_error` | 是否为错误 | +| `result.async` | 是否异步执行 | +| `result.media` | 媒体引用列表 | +| `result.artifact_tags` | 本地产物路径标签 | +| `result.response_handled` | 是否已处理响应 | +| `duration` | 执行耗时（纳秒） | + +### 响应 + +```json +{ + "jsonrpc": "2.0", + "id": 5, + "result": { + "action": "continue" + } +} +``` + +--- + +## 6. `hook.approve_tool` + +审批型 hook，用于决定是否允许执行敏感工具。 + +### 请求 + +```json +{ + "jsonrpc": "2.0", + "id": 6, + "method": "hook.approve_tool", + "params": { + "meta": { + "AgentID": "agent-1", + "TurnID": "turn-1", + "SessionKey": "session-1" + }, + "tool": "bash", + "arguments": { + "command": "rm -rf /" + }, + "channel": "cli", + "chat_id": "chat-1" + } +} +``` + +### 响应（批准） + +```json +{ + "jsonrpc": "2.0", + "id": 6, + "result": { + "approved": true + } +} +``` + +### 响应（拒绝） + +```json +{ + "jsonrpc": "2.0", + "id": 6, + "result": { + "approved": false, + "reason": "危险命令，禁止执行" + } +} +``` + +--- + +## 7. `hook.event`（notification） + +观察型事件，仅广播，无需响应。`id` 为 `0` 或不存在。 + +```json +{ + "jsonrpc": "2.0", + "method": "hook.event", + "params": { + "Kind": "tool_exec_start", + "Meta": { + "AgentID": "agent-1", + "TurnID": "turn-1" + }, + "Payload": { + "Tool": "echo_text", + "Arguments": {"text": "hello"} + } + } +} +``` + +常见 `Kind` 值： +- `turn_start` / `turn_end` +- `llm_request` / `llm_response` +- `tool_exec_start` / `tool_exec_end` / `tool_exec_skipped` +- `steering_injected` +- `interrupt_received` +- `error` + +--- + +## action 可选值 + +| action | 适用 hook | 效果 | +|--------|----------|------| +| `continue` | 所有拦截型 | 放行，不做修改 | +| `modify` | `before_llm`, `before_tool`, `after_llm`, `after_tool` | 改写请求/响应后放行 | +| `respond` | `before_tool` | 直接返回工具结果，跳过实际执行 | +| `deny_tool` | `before_tool` | 拒绝执行该工具 | +| `abort_turn` | 所有拦截型 | 中止当前 turn，返回错误 | +| `hard_abort` | 所有拦截型 | 强制终止整个 agent loop | + +--- + +## 完整流程示例 + +```json +{"jsonrpc":"2.0","id":1,"method":"hook.hello","params":{"name":"my_hook","version":1,"modes":["tool","approve"]}} +{"jsonrpc":"2.0","id":1,"result":{"ok":true,"name":"my_hook"}} +{"jsonrpc":"2.0","id":2,"method":"hook.before_llm","params":{"model":"claude-sonnet","messages":[{"role":"user","content":"hello"}],"tools":[]}} +{"jsonrpc":"2.0","id":2,"result":{"action":"continue"}} +{"jsonrpc":"2.0","id":3,"method":"hook.before_tool","params":{"tool":"bash","arguments":{"command":"ls"}}} +{"jsonrpc":"2.0","id":3,"result":{"action":"continue"}} +{"jsonrpc":"2.0","id":4,"method":"hook.approve_tool","params":{"tool":"bash","arguments":{"command":"ls"}}} +{"jsonrpc":"2.0","id":4,"result":{"approved":true}} +{"jsonrpc":"2.0","id":5,"method":"hook.after_tool","params":{"tool":"bash","arguments":{"command":"ls"},"result":{"for_llm":"file1.txt\nfile2.txt"},"duration":5000000}} +{"jsonrpc":"2.0","id":5,"result":{"action":"continue"}} +{"jsonrpc":"2.0","id":6,"method":"hook.after_llm","params":{"model":"claude-sonnet","response":{"role":"assistant","content":"已列出文件"}}} +{"jsonrpc":"2.0","id":6,"result":{"action":"continue"}} +``` + +--- + +## 通过 `before_llm` 和 `before_tool` 实现插件工具注入 + +插件工具注入的标准流程： + +1. 在 `before_llm` 中注入工具定义，让 LLM 知道有这个工具可用 +2. 在 `before_tool` 中使用 `respond` action 直接返回工具执行结果 + +### `before_llm` 注入工具定义 + +```python +def handle_before_llm(params: dict) -> dict: + tools = params.get("tools", []) + + # 添加插件工具定义 + tools.append({ + "type": "function", + "function": { + "name": "my_plugin_tool", + "description": "插件提供的工具", + "parameters": { + "type": "object", + "properties": { + "input": {"type": "string", "description": "输入内容"} + }, + "required": ["input"] + } + } + }) + + return { + "action": "modify", + "request": { + "model": params["model"], + "messages": params["messages"], + "tools": tools, + "options": params.get("options", {}) + } + } +``` + +### `before_tool` 返回执行结果 + +```python +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + + if tool == "my_plugin_tool": + # 在这里实现工具逻辑 + args = params.get("arguments", {}) + input_text = args.get("input", "") + + # 直接返回结果，无需在 ToolRegistry 注册 + return { + "action": "respond", + "result": { + "for_llm": f"插件工具执行成功，输入: {input_text}", + "silent": False, + "is_error": False + } + } + + return {"action": "continue"} +``` + +通过这种方式，外部 hook 可以完全实现插件工具，无需在 PicoClaw 内部注册任何工具实现。 \ No newline at end of file diff --git a/docs/hooks/plugin-tool-injection.md b/docs/hooks/plugin-tool-injection.md new file mode 100644 index 000000000..9e699867b --- /dev/null +++ b/docs/hooks/plugin-tool-injection.md @@ -0,0 +1,587 @@ +# Plugin Tool Injection Example + +This document demonstrates how to use PicoClaw's hook system to implement external plugin tool injection, allowing LLM to call tools implemented by external hook processes. + +--- + +## Core Principle + +Through the hook system's `respond` action, external hooks can: + +1. Inject tool **definitions** in `before_llm`, letting LLM know the tool is available +2. Return tool **execution results** directly in `before_tool` using `respond` action, skipping ToolRegistry + +This way, external hooks can fully implement plugin tools without registering any tools inside PicoClaw. + +--- + +## Complete Example: Weather Query Plugin + +Below is a complete Python hook example implementing a weather query plugin tool. + +### 1. Hook Script Implementation + +Save as `/tmp/weather_plugin.py`: + +```python +#!/usr/bin/env python3 +"""Weather query plugin hook example""" +from __future__ import annotations + +import json +import sys +import signal +from typing import Any + +# Simulated weather data +WEATHER_DATA = { + "Beijing": {"temp": 15, "weather": "Sunny", "humidity": 45}, + "Shanghai": {"temp": 18, "weather": "Cloudy", "humidity": 60}, + "Guangzhou": {"temp": 25, "weather": "Sunny", "humidity": 70}, + "Shenzhen": {"temp": 26, "weather": "Cloudy", "humidity": 75}, +} + + +def get_weather(city: str) -> dict: + """Get weather data (simulated)""" + data = WEATHER_DATA.get(city) + if data: + return { + "for_llm": f"{city} weather: {data['weather']}, temperature {data['temp']}°C, humidity {data['humidity']}%", + "for_user": "", + "silent": False, + "is_error": False, + } + return { + "for_llm": f"Weather data not found for city {city}", + "for_user": "", + "silent": False, + "is_error": True, + } + + +def handle_hello(params: dict) -> dict: + return {"ok": True, "name": "weather-plugin"} + + +def handle_before_llm(params: dict) -> dict: + """Inject weather query tool definition""" + tools = params.get("tools", []) + + # Add weather query tool + tools.append({ + "type": "function", + "function": { + "name": "get_weather", + "description": "Query weather information for a specified city", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string", + "description": "City name, e.g.: Beijing, Shanghai, Guangzhou" + } + }, + "required": ["city"] + } + } + }) + + return { + "action": "modify", + "request": { + "model": params.get("model"), + "messages": params.get("messages", []), + "tools": tools, + "options": params.get("options", {}), + } + } + + +def handle_before_tool(params: dict) -> dict: + """Handle tool call, return result directly""" + tool = params.get("tool", "") + args = params.get("arguments", {}) + + if tool == "get_weather": + city = args.get("city", "") + result = get_weather(city) + + # Use respond action to return result directly, skip ToolRegistry + return { + "action": "respond", + "result": result, + } + + # Other tools continue normal flow + return {"action": "continue"} + + +def handle_request(method: str, params: dict) -> dict: + if method == "hook.hello": + return handle_hello(params) + if method == "hook.before_llm": + return handle_before_llm(params) + if method == "hook.before_tool": + return handle_before_tool(params) + if method == "hook.after_llm": + return {"action": "continue"} + if method == "hook.after_tool": + return {"action": "continue"} + if method == "hook.approve_tool": + return {"approved": True} + raise KeyError(f"method not found: {method}") + + +def send_response(message_id: int, result: Any | None = None, error: str | None = None) -> None: + payload: dict[str, Any] = { + "jsonrpc": "2.0", + "id": message_id, + } + if error is not None: + payload["error"] = {"code": -32000, "message": error} + else: + payload["result"] = result if result is not None else {} + + sys.stdout.write(json.dumps(payload, ensure_ascii=True) + "\n") + sys.stdout.flush() + + +def main() -> int: + for raw_line in sys.stdin: + line = raw_line.strip() + if not line: + continue + + try: + message = json.loads(line) + except json.JSONDecodeError: + continue + + method = message.get("method") + message_id = message.get("id", 0) + params = message.get("params") or {} + + if not message_id: + continue + + try: + result = handle_request(str(method or ""), params) + send_response(int(message_id), result=result) + except KeyError as exc: + send_response(int(message_id), error=str(exc)) + except Exception as exc: + send_response(int(message_id), error=f"unexpected error: {exc}") + + return 0 + + +if __name__ == "__main__": + signal.signal(signal.SIGINT, lambda *_: raise SystemExit(0)) + signal.signal(signal.SIGTERM, lambda *_: raise SystemExit(0)) + raise SystemExit(main()) +``` + +### 2. Configure PicoClaw + +Add hook configuration in the config file: + +```json +{ + "hooks": { + "enabled": true, + "processes": { + "weather_plugin": { + "enabled": true, + "priority": 100, + "transport": "stdio", + "command": ["python3", "/tmp/weather_plugin.py"], + "intercept": ["before_llm", "before_tool"] + } + } + } +} +``` + +### 3. Test Results + +When user asks "What's the weather in Beijing today?": + +1. PicoClaw sends `hook.before_llm`, hook injects `get_weather` tool definition +2. LLM sees tool definition, decides to call `get_weather(city="Beijing")` +3. PicoClaw sends `hook.before_tool`, hook uses `respond` action to return weather data +4. LLM receives result, replies to user "Beijing is sunny today, temperature 15°C" + +--- + +## Flow Diagram + +``` +User: "What's the weather in Beijing today?" + ↓ + PicoClaw + ↓ + hook.before_llm + ↓ (inject get_weather tool definition) + LLM request + ↓ + LLM decides to call get_weather(city="Beijing") + ↓ + hook.before_tool + ↓ (respond action returns weather data) + Return result directly to LLM + ↓ (skip ToolRegistry) + LLM replies: "Beijing is sunny today, temperature 15°C" +``` + +--- + +## Key Points + +### `before_llm` Inject Tool Definition + +Tool definition follows OpenAI function calling format: + +```json +{ + "type": "function", + "function": { + "name": "tool_name", + "description": "tool description", + "parameters": { + "type": "object", + "properties": { + "param_name": { + "type": "string", + "description": "parameter description" + } + }, + "required": ["list of required parameters"] + } + } +} +``` + +### `before_tool` Use respond Action + +`respond` action response format: + +```json +{ + "action": "respond", + "result": { + "for_llm": "Content returned to LLM", + "for_user": "Optional, content sent to user", + "silent": false, + "is_error": false, + "media": ["Optional, media reference list"], + "response_handled": false + } +} +``` + +| Field | Description | +|-------|-------------| +| `for_llm` | Required, LLM will see this content | +| `for_user` | Optional, sent directly to user | +| `silent` | When true, not sent to user | +| `is_error` | When true, indicates execution failure | +| `media` | Optional, media file references (images, files, etc.) | +| `response_handled` | When true, indicates user request is handled, turn will end | + +--- + +## Media File Handling + +The `respond` action supports returning media files (images, files, etc.). There are two processing modes: + +### 1. Automatic Delivery (`response_handled=true`) + +When `response_handled=true`, media files are automatically sent to the user and the turn ends: + +```json +{ + "action": "respond", + "result": { + "for_llm": "Image sent to user", + "for_user": "", + "media": ["media://abc123"], + "response_handled": true + } +} +``` + +Use cases: +- Image generation plugin directly returning results +- File download plugin sending files to user + +### 2. LLM Visible (`response_handled=false`) + +When `response_handled=false`, media references are passed to the LLM, which can see the content in the next request: + +```json +{ + "action": "respond", + "result": { + "for_llm": "Image loaded, path: /tmp/image.png [file:/tmp/image.png]", + "media": ["media://abc123"] + } +} +``` + +After seeing the content, the LLM can decide: +- Use `send_file` tool to send to user +- Analyze image content and reply to user +- Other processing approaches + +### Media Reference Format + +Media references use the `media://` protocol: + +``` +media:// +``` + +These references are managed by PicoClaw's MediaStore and can be: +- Sent to user via channel +- Converted to base64 in LLM vision requests + +### Alternative: Use Existing Tools + +If the plugin generates files, you can return the file path and let the LLM call `send_file` or similar tools: + +```json +{ + "action": "respond", + "result": { + "for_llm": "Image generated, saved at /tmp/generated_image.png. Use send_file tool to send to user.", + "for_user": "", + "silent": false + } +} +``` + +This approach: +- More decoupled, LLM decides when to send +- Leverages existing tool mechanisms +- Supports batch sending, delayed sending, etc. + +--- + +## Multi-Tool Injection Example + +Multiple tools can be injected simultaneously: + +```python +def handle_before_llm(params: dict) -> dict: + tools = params.get("tools", []) + + # Tool 1: Weather query + tools.append({ + "type": "function", + "function": { + "name": "get_weather", + "description": "Query city weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"] + } + } + }) + + # Tool 2: Calculator + tools.append({ + "type": "function", + "function": { + "name": "calculate", + "description": "Perform mathematical calculations", + "parameters": { + "type": "object", + "properties": { + "expression": {"type": "string", "description": "Mathematical expression"} + }, + "required": ["expression"] + } + } + }) + + return { + "action": "modify", + "request": { + "model": params.get("model"), + "messages": params.get("messages", []), + "tools": tools, + "options": params.get("options", {}), + } + } + + +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + args = params.get("arguments", {}) + + if tool == "get_weather": + return { + "action": "respond", + "result": get_weather(args.get("city", "")), + } + + if tool == "calculate": + # Simple calculation example + try: + expr = args.get("expression", "") + result = eval(expr) # Note: needs security handling in actual use + return { + "action": "respond", + "result": { + "for_llm": f"Calculation result: {result}", + "silent": False, + "is_error": False, + }, + } + except Exception as e: + return { + "action": "respond", + "result": { + "for_llm": f"Calculation error: {e}", + "silent": False, + "is_error": True, + }, + } + + return {"action": "continue"} +``` + +--- + +## Coexistence with Built-in Tools + +Injected plugin tools coexist with PicoClaw built-in tools: + +- Built-in tools (like `bash`, `read_file`) execute normally through ToolRegistry +- Plugin tools return results through hook's `respond` action +- `handle_before_tool` only handles plugin tools, other tools return `continue` + +--- + +## Go In-Process Hook Example + +If you need to implement plugin tool injection in Go code: + +```go +package myhooks + +import ( + "context" + "github.com/sipeed/picoclaw/pkg/agent" + "github.com/sipeed/picoclaw/pkg/tools" +) + +type WeatherPluginHook struct{} + +func (h *WeatherPluginHook) BeforeLLM( + ctx context.Context, + req *agent.LLMHookRequest, +) (*agent.LLMHookRequest, agent.HookDecision, error) { + // Inject tool definition + req.Tools = append(req.Tools, agent.ToolDefinition{ + Type: "function", + Function: agent.FunctionDefinition{ + Name: "get_weather", + Description: "Query city weather", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "city": map[string]any{ + "type": "string", + "description": "City name", + }, + }, + "required": []string{"city"}, + }, + }, + }) + + return req, agent.HookDecision{Action: agent.HookActionContinue}, nil +} + +func (h *WeatherPluginHook) BeforeTool( + ctx context.Context, + call *agent.ToolCallHookRequest, +) (*agent.ToolCallHookRequest, agent.HookDecision, error) { + if call.Tool == "get_weather" { + city := call.Arguments["city"].(string) + + // Set HookResult, use respond action + next := call.Clone() + next.HookResult = &tools.ToolResult{ + ForLLM: getWeatherData(city), + Silent: false, + IsError: false, + } + + return next, agent.HookDecision{Action: agent.HookActionRespond}, nil + } + + return call, agent.HookDecision{Action: agent.HookActionContinue}, nil +} + +func getWeatherData(city string) string { + // Implement weather query logic + return fmt.Sprintf("%s weather: Sunny, temperature 20°C", city) +} +``` + +--- + +## Summary + +Through the hook system's `respond` action, external processes can: + +1. **Inject tool definitions**: Let LLM know new tools are available +2. **Provide tool implementation**: Return execution results directly, no need to register in ToolRegistry +3. **Coexist with built-in tools**: Does not affect normal operation of PicoClaw's original tools + +This provides a flexible and elegant solution for plugin development. + +--- + +## Security Boundaries + +### Bypassing Approval Checks + +**Important**: The `respond` action bypasses `ApproveTool` approval checks. + +This means: +- A `before_tool` hook can return `respond` for **any tool name**, including sensitive tools (like `bash`) +- The tool won't go through the approval process, directly returning the hook-provided result +- This is designed for plugin tools but introduces security risks + +### Security Recommendations + +1. **Review hook configuration**: Ensure only trusted hook processes are enabled +2. **Limit hook scope**: Add your own security checks in hook implementation +3. **Use `deny_tool` for rejection**: Use `deny_tool` action instead of `respond` with error for denying execution + +### Example: Hook-Internal Security Check + +```python +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + args = params.get("arguments", {}) + + # Security check: only handle plugin tools + if tool in ["get_weather", "calculate"]: + return { + "action": "respond", + "result": execute_plugin_tool(tool, args), + } + + # Other tools continue normal flow (will go through approval) + return {"action": "continue"} +``` + +This ensures the hook only affects plugin tools, not system tool approval flow. \ No newline at end of file diff --git a/docs/hooks/plugin-tool-injection.zh.md b/docs/hooks/plugin-tool-injection.zh.md new file mode 100644 index 000000000..ccc7ff7f6 --- /dev/null +++ b/docs/hooks/plugin-tool-injection.zh.md @@ -0,0 +1,587 @@ +# 插件工具注入示例 + +本文档展示如何利用 PicoClaw 的 hook 系统实现外部插件工具注入，让 LLM 能调用由外部 hook 进程实现的工具。 + +--- + +## 核心原理 + +通过 hook 系统的 `respond` action，外部 hook 可以： + +1. 在 `before_llm` 中注入工具**定义**，让 LLM 知道有这个工具可用 +2. 在 `before_tool` 中使用 `respond` action 直接返回工具**执行结果**，跳过 ToolRegistry + +这样，外部 hook 可以完全实现插件工具，无需在 PicoClaw 内部注册任何工具。 + +--- + +## 完整示例：天气查询插件 + +下面是一个完整的 Python hook 示例，实现一个天气查询插件工具。 + +### 1. Hook 脚本实现 + +保存为 `/tmp/weather_plugin.py`： + +```python +#!/usr/bin/env python3 +"""天气查询插件 hook 示例""" +from __future__ import annotations + +import json +import sys +import signal +from typing import Any + +# 模拟天气数据 +WEATHER_DATA = { + "北京": {"temp": 15, "weather": "晴", "humidity": 45}, + "上海": {"temp": 18, "weather": "多云", "humidity": 60}, + "广州": {"temp": 25, "weather": "晴", "humidity": 70}, + "深圳": {"temp": 26, "weather": "多云", "humidity": 75}, +} + + +def get_weather(city: str) -> dict: + """获取天气数据（模拟）""" + data = WEATHER_DATA.get(city) + if data: + return { + "for_llm": f"{city}天气：{data['weather']}，温度{data['temp']}°C，湿度{data['humidity']}%", + "for_user": "", + "silent": False, + "is_error": False, + } + return { + "for_llm": f"未找到城市 {city} 的天气数据", + "for_user": "", + "silent": False, + "is_error": True, + } + + +def handle_hello(params: dict) -> dict: + return {"ok": True, "name": "weather-plugin"} + + +def handle_before_llm(params: dict) -> dict: + """注入天气查询工具定义""" + tools = params.get("tools", []) + + # 添加天气查询工具 + tools.append({ + "type": "function", + "function": { + "name": "get_weather", + "description": "查询指定城市的天气信息", + "parameters": { + "type": "object", + "properties": { + "city": { + "type": "string", + "description": "城市名称，如：北京、上海、广州" + } + }, + "required": ["city"] + } + } + }) + + return { + "action": "modify", + "request": { + "model": params.get("model"), + "messages": params.get("messages", []), + "tools": tools, + "options": params.get("options", {}), + } + } + + +def handle_before_tool(params: dict) -> dict: + """处理工具调用，直接返回结果""" + tool = params.get("tool", "") + args = params.get("arguments", {}) + + if tool == "get_weather": + city = args.get("city", "") + result = get_weather(city) + + # 使用 respond action 直接返回结果，跳过 ToolRegistry + return { + "action": "respond", + "result": result, + } + + # 其他工具继续正常流程 + return {"action": "continue"} + + +def handle_request(method: str, params: dict) -> dict: + if method == "hook.hello": + return handle_hello(params) + if method == "hook.before_llm": + return handle_before_llm(params) + if method == "hook.before_tool": + return handle_before_tool(params) + if method == "hook.after_llm": + return {"action": "continue"} + if method == "hook.after_tool": + return {"action": "continue"} + if method == "hook.approve_tool": + return {"approved": True} + raise KeyError(f"method not found: {method}") + + +def send_response(message_id: int, result: Any | None = None, error: str | None = None) -> None: + payload: dict[str, Any] = { + "jsonrpc": "2.0", + "id": message_id, + } + if error is not None: + payload["error"] = {"code": -32000, "message": error} + else: + payload["result"] = result if result is not None else {} + + sys.stdout.write(json.dumps(payload, ensure_ascii=True) + "\n") + sys.stdout.flush() + + +def main() -> int: + for raw_line in sys.stdin: + line = raw_line.strip() + if not line: + continue + + try: + message = json.loads(line) + except json.JSONDecodeError: + continue + + method = message.get("method") + message_id = message.get("id", 0) + params = message.get("params") or {} + + if not message_id: + continue + + try: + result = handle_request(str(method or ""), params) + send_response(int(message_id), result=result) + except KeyError as exc: + send_response(int(message_id), error=str(exc)) + except Exception as exc: + send_response(int(message_id), error=f"unexpected error: {exc}") + + return 0 + + +if __name__ == "__main__": + signal.signal(signal.SIGINT, lambda *_: raise SystemExit(0)) + signal.signal(signal.SIGTERM, lambda *_: raise SystemExit(0)) + raise SystemExit(main()) +``` + +### 2. 配置 PicoClaw + +在配置文件中添加 hook 配置： + +```json +{ + "hooks": { + "enabled": true, + "processes": { + "weather_plugin": { + "enabled": true, + "priority": 100, + "transport": "stdio", + "command": ["python3", "/tmp/weather_plugin.py"], + "intercept": ["before_llm", "before_tool"] + } + } + } +} +``` + +### 3. 测试效果 + +当用户问"北京今天天气怎么样？"时： + +1. PicoClaw 发送 `hook.before_llm`，hook 注入 `get_weather` 工具定义 +2. LLM 看到工具定义，决定调用 `get_weather(city="北京")` +3. PicoClaw 发送 `hook.before_tool`，hook 使用 `respond` action 返回天气数据 +4. LLM 收到结果，回复用户"北京今天晴天，温度15°C" + +--- + +## 流程图解 + +``` +用户: "北京今天天气怎么样？" + ↓ + PicoClaw + ↓ + hook.before_llm + ↓ (注入 get_weather 工具定义) + LLM 请求 + ↓ + LLM 决定调用 get_weather(city="北京") + ↓ + hook.before_tool + ↓ (respond action 返回天气数据) + 直接返回结果给 LLM + ↓ (跳过 ToolRegistry) + LLM 回复: "北京今天晴天，温度15°C" +``` + +--- + +## 关键点说明 + +### `before_llm` 注入工具定义 + +工具定义遵循 OpenAI function calling 格式： + +```json +{ + "type": "function", + "function": { + "name": "工具名称", + "description": "工具描述", + "parameters": { + "type": "object", + "properties": { + "参数名": { + "type": "string", + "description": "参数描述" + } + }, + "required": ["必需参数列表"] + } + } +} +``` + +### `before_tool` 使用 respond action + +`respond` action 的响应格式： + +```json +{ + "action": "respond", + "result": { + "for_llm": "返回给 LLM 的内容", + "for_user": "可选，发送给用户的内容", + "silent": false, + "is_error": false, + "media": ["可选，媒体引用列表"], + "response_handled": false + } +} +``` + +| 字段 | 说明 | +|------|------| +| `for_llm` | 必须，LLM 会看到这个内容 | +| `for_user` | 可选，直接发送给用户 | +| `silent` | 为 true 时不发送给用户 | +| `is_error` | 为 true 时表示执行失败 | +| `media` | 可选，媒体文件引用列表（如图片、文件） | +| `response_handled` | 为 true 时表示已处理用户请求，轮次将结束 | + +--- + +## 媒体文件处理 + +`respond` action 支持返回媒体文件（图片、文件等）。有两种处理方式： + +### 1. 自动发送（`response_handled=true`） + +当 `response_handled=true` 时，媒体文件会自动发送给用户，轮次结束： + +```json +{ + "action": "respond", + "result": { + "for_llm": "图片已发送给用户", + "for_user": "", + "media": ["media://abc123"], + "response_handled": true + } +} +``` + +适用场景： +- 图像生成插件直接返回结果 +- 文件下载插件发送文件给用户 + +### 2. LLM 可见（`response_handled=false`） + +当 `response_handled=false` 时，媒体引用会传递给 LLM，LLM 可以在下一轮请求中看到内容： + +```json +{ + "action": "respond", + "result": { + "for_llm": "图片已加载，路径：/tmp/image.png [file:/tmp/image.png]", + "media": ["media://abc123"] + } +} +``` + +LLM 看到内容后，可以自主决定： +- 使用 `send_file` 工具发送给用户 +- 分析图片内容并回复用户 +- 其他处理方式 + +### 媒体引用格式 + +媒体引用使用 `media://` 协议： + +``` +media:// +``` + +这些引用由 PicoClaw 的 MediaStore 管理，可以： +- 通过 channel 发送给用户 +- 在 LLM vision 请求中转换为 base64 + +### 替代方案：使用现有工具 + +如果插件生成文件，可以返回文件路径让 LLM 调用 `send_file` 等工具： + +```json +{ + "action": "respond", + "result": { + "for_llm": "图片已生成，保存在 /tmp/generated_image.png。使用 send_file 工具发送给用户。", + "for_user": "", + "silent": false + } +} +``` + +这种方式： +- 更解耦，LLM 自主决策发送时机 +- 利用现有工具机制 +- 支持批量发送、延迟发送等场景 + +--- + +## 多工具注入示例 + +可以同时注入多个工具： + +```python +def handle_before_llm(params: dict) -> dict: + tools = params.get("tools", []) + + # 工具1：天气查询 + tools.append({ + "type": "function", + "function": { + "name": "get_weather", + "description": "查询城市天气", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "城市名称"} + }, + "required": ["city"] + } + } + }) + + # 工具2：计算器 + tools.append({ + "type": "function", + "function": { + "name": "calculate", + "description": "执行数学计算", + "parameters": { + "type": "object", + "properties": { + "expression": {"type": "string", "description": "数学表达式"} + }, + "required": ["expression"] + } + } + }) + + return { + "action": "modify", + "request": { + "model": params.get("model"), + "messages": params.get("messages", []), + "tools": tools, + "options": params.get("options", {}), + } + } + + +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + args = params.get("arguments", {}) + + if tool == "get_weather": + return { + "action": "respond", + "result": get_weather(args.get("city", "")), + } + + if tool == "calculate": + # 简单计算示例 + try: + expr = args.get("expression", "") + result = eval(expr) # 注意：实际使用时需要安全处理 + return { + "action": "respond", + "result": { + "for_llm": f"计算结果: {result}", + "silent": False, + "is_error": False, + }, + } + except Exception as e: + return { + "action": "respond", + "result": { + "for_llm": f"计算错误: {e}", + "silent": False, + "is_error": True, + }, + } + + return {"action": "continue"} +``` + +--- + +## 与内置工具共存 + +注入的插件工具与 PicoClaw 内置工具共存： + +- 内置工具（如 `bash`、`read_file`）正常通过 ToolRegistry 执行 +- 插件工具通过 hook 的 `respond` action 返回结果 +- `handle_before_tool` 中只处理插件工具，其他工具返回 `continue` + +--- + +## Go 进程内 Hook 示例 + +如果需要在 Go 代码中实现插件工具注入： + +```go +package myhooks + +import ( + "context" + "github.com/sipeed/picoclaw/pkg/agent" + "github.com/sipeed/picoclaw/pkg/tools" +) + +type WeatherPluginHook struct{} + +func (h *WeatherPluginHook) BeforeLLM( + ctx context.Context, + req *agent.LLMHookRequest, +) (*agent.LLMHookRequest, agent.HookDecision, error) { + // 注入工具定义 + req.Tools = append(req.Tools, agent.ToolDefinition{ + Type: "function", + Function: agent.FunctionDefinition{ + Name: "get_weather", + Description: "查询城市天气", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "city": map[string]any{ + "type": "string", + "description": "城市名称", + }, + }, + "required": []string{"city"}, + }, + }, + }) + + return req, agent.HookDecision{Action: agent.HookActionContinue}, nil +} + +func (h *WeatherPluginHook) BeforeTool( + ctx context.Context, + call *agent.ToolCallHookRequest, +) (*agent.ToolCallHookRequest, agent.HookDecision, error) { + if call.Tool == "get_weather" { + city := call.Arguments["city"].(string) + + // 设置 HookResult，使用 respond action + next := call.Clone() + next.HookResult = &tools.ToolResult{ + ForLLM: getWeatherData(city), + Silent: false, + IsError: false, + } + + return next, agent.HookDecision{Action: agent.HookActionRespond}, nil + } + + return call, agent.HookDecision{Action: agent.HookActionContinue}, nil +} + +func getWeatherData(city string) string { + // 实现天气查询逻辑 + return fmt.Sprintf("%s天气：晴，温度20°C", city) +} +``` + +--- + +## 总结 + +通过 hook 系统的 `respond` action，外部进程可以： + +1. **注入工具定义**：让 LLM 知道有新工具可用 +2. **提供工具实现**：直接返回执行结果，无需注册到 ToolRegistry +3. **与内置工具共存**：不影响 PicoClaw 原有工具的正常运行 + +这为插件开发提供了灵活、优雅的解决方案。 + +--- + +## 安全边界说明 + +### 绕过审批检查 + +**重要**：`respond` action 会绕过 `ApproveTool` 审批检查。 + +这意味着： +- `before_tool` hook 可以为**任何工具名称**返回 `respond`，包括敏感工具（如 `bash`） +- 工具不会经过审批流程，直接返回 hook 提供的结果 +- 这是为了支持插件工具而设计，但也带来了安全风险 + +### 安全建议 + +1. **审查 hook 配置**：确保只有可信的 hook 进程被启用 +2. **限制 hook 权限**：在 hook 实现中添加自己的安全检查 +3. **优先使用 `deny_tool`**：对于拒绝执行，使用 `deny_tool` action 而非 `respond` 返回错误 + +### 示例：hook 内置安全检查 + +```python +def handle_before_tool(params: dict) -> dict: + tool = params.get("tool", "") + args = params.get("arguments", {}) + + # 安全检查：只处理插件工具 + if tool in ["get_weather", "calculate"]: + return { + "action": "respond", + "result": execute_plugin_tool(tool, args), + } + + # 其他工具继续正常流程（会经过审批） + return {"action": "continue"} +``` + +这样可以确保 hook 只影响插件工具，不影响系统工具的审批流程。 \ No newline at end of file diff --git a/docs/providers.md b/docs/providers.md index 9bb95446c..d03fbab3e 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -122,6 +122,7 @@ This design also enables **multi-agent support** with flexible provider selectio | `max_tokens_field` | string | No | Override the max tokens field name in request body (e.g., `max_completion_tokens` for o1 models) | | `thinking_level` | string | No | Extended thinking level: `off`, `low`, `medium`, `high`, `xhigh`, or `adaptive` | | `extra_body` | object | No | Additional fields to inject into every request body | +| `custom_headers` | object | No | Additional HTTP headers to inject into every request (e.g., `{"X-Source":"coding-plan"}`). If a key matches a built-in header, the custom value overrides the built-in one (e.g., `Authorization`, `User-Agent`, `Content-Type`, `Accept`). | | `rpm` | int | No | Per-minute request rate limit | | `fallbacks` | string[] | No | Fallback model names for automatic failover | | `enabled` | bool | No | Whether this model entry is active (default: `true`) | diff --git a/docs/zh/providers.md b/docs/zh/providers.md index 6048b929f..7b3930f6f 100644 --- a/docs/zh/providers.md +++ b/docs/zh/providers.md @@ -118,6 +118,7 @@ | `max_tokens_field` | string | 否 | 覆盖请求体中 max tokens 的字段名（如 o1 模型使用 `max_completion_tokens`） | | `thinking_level` | string | 否 | 扩展思考级别：`off`、`low`、`medium`、`high`、`xhigh` 或 `adaptive` | | `extra_body` | object | 否 | 注入到每个请求体中的额外字段 | +| `custom_headers` | object | 否 | 注入到每个请求中的额外 HTTP 请求头（例如 `{"X-Source":"coding-plan"}`）。若键名与内置请求头同名，会覆盖内置值（如 `Authorization`、`User-Agent`、`Content-Type`、`Accept`）。 | | `rpm` | int | 否 | 每分钟请求速率限制 | | `fallbacks` | string[] | 否 | 自动故障转移的备用模型名称 | | `enabled` | bool | 否 | 是否启用此模型条目（默认：`true`） | diff --git a/go.mod b/go.mod index 7044e29e8..1ff7cb306 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/sipeed/picoclaw -go 1.25.8 +go 1.25.9 require ( fyne.io/systray v1.12.0 @@ -8,6 +8,7 @@ require ( github.com/SevereCloud/vksdk/v3 v3.3.1 github.com/adhocore/gronx v1.19.6 github.com/anthropics/anthropic-sdk-go v1.26.0 + github.com/atc0005/go-teams-notify/v2 v2.14.0 github.com/atotto/clipboard v0.1.4 github.com/aws/aws-sdk-go-v2 v1.41.5 github.com/aws/aws-sdk-go-v2/config v1.32.12 @@ -29,7 +30,7 @@ require ( github.com/mymmrac/telego v1.7.0 github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1 github.com/openai/openai-go/v3 v3.22.0 - github.com/pion/rtp v1.8.7 + github.com/pion/rtp v1.10.1 github.com/pion/webrtc/v3 v3.3.6 github.com/rivo/tview v0.42.0 github.com/rs/zerolog v1.35.0 @@ -45,7 +46,7 @@ require ( google.golang.org/protobuf v1.36.11 gopkg.in/yaml.v3 v3.0.1 maunium.net/go/mautrix v0.26.4 - modernc.org/sqlite v1.47.0 + modernc.org/sqlite v1.48.0 rsc.io/qr v0.2.0 ) @@ -84,7 +85,6 @@ require ( github.com/petermattis/goid v0.0.0-20260226131333-17d1149c6ac6 // indirect github.com/pion/randutil v0.1.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/reiver/go-porterstemmer v1.0.1 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/rivo/uniseg v0.4.7 // indirect github.com/segmentio/asm v1.1.3 // indirect diff --git a/go.sum b/go.sum index 67ccd2b12..765a3211a 100644 --- a/go.sum +++ b/go.sum @@ -21,6 +21,8 @@ github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwTo github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/anthropics/anthropic-sdk-go v1.26.0 h1:oUTzFaUpAevfuELAP1sjL6CQJ9HHAfT7CoSYSac11PY= github.com/anthropics/anthropic-sdk-go v1.26.0/go.mod h1:qUKmaW+uuPB64iy1l+4kOSvaLqPXnHTTBKH6RVZ7q5Q= +github.com/atc0005/go-teams-notify/v2 v2.14.0 h1:7N+xw+COnYANLREaAveQ65rsNQ12nIZJED9nMLyscCo= +github.com/atc0005/go-teams-notify/v2 v2.14.0/go.mod h1:EECsWM2b0Hvoz7O+QdlsvyN2KCUOFQCGj8bUBXv3A3Q= github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= github.com/aws/aws-sdk-go-v2 v1.41.5 h1:dj5kopbwUsVUVFgO4Fi5BIT3t4WyqIDjGKCangnV/yY= @@ -207,15 +209,13 @@ github.com/petermattis/goid v0.0.0-20260226131333-17d1149c6ac6 h1:rh2lKw/P/EqHa7 github.com/petermattis/goid v0.0.0-20260226131333-17d1149c6ac6/go.mod h1:pxMtw7cyUw6B2bRH0ZBANSPg+AoSud1I1iyJHI69jH4= github.com/pion/randutil v0.1.0 h1:CFG1UdESneORglEsnimhUjf33Rwjubwj6xfiOXBa3mA= github.com/pion/randutil v0.1.0/go.mod h1:XcJrSMMbbMRhASFVOlj/5hQial/Y8oH/HVo7TBZq+j8= -github.com/pion/rtp v1.8.7 h1:qslKkG8qxvQ7hqaxkmL7Pl0XcUm+/Er7nMnu6Vq+ZxM= -github.com/pion/rtp v1.8.7/go.mod h1:pBGHaFt/yW7bf1jjWAoUjpSNoDnw98KTMg+jWWvziqU= +github.com/pion/rtp v1.10.1 h1:xP1prZcCTUuhO2c83XtxyOHJteISg6o8iPsE2acaMtA= +github.com/pion/rtp v1.10.1/go.mod h1:rF5nS1GqbR7H/TCpKwylzeq6yDM+MM6k+On5EgeThEM= github.com/pion/webrtc/v3 v3.3.6 h1:7XAh4RPtlY1Vul6/GmZrv7z+NnxKA6If0KStXBI2ZLE= github.com/pion/webrtc/v3 v3.3.6/go.mod h1:zyN7th4mZpV27eXybfR/cnUf3J2DRy8zw/mdjD9JTNM= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= 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/reiver/go-porterstemmer v1.0.1 h1:WyERBkASXgoXrTwq/IQ6wyNj/YG7j/ZURvTuMCoud5w= -github.com/reiver/go-porterstemmer v1.0.1/go.mod h1:Z8uL/f/7UEwaeAJNwx1sO8kbqXiEuQieNuD735hLrSU= 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/tview v0.42.0 h1:b/ftp+RxtDsHSaynXTbJb+/n/BxDEi+W3UfF5jILK6c= @@ -458,8 +458,8 @@ modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= modernc.org/opt v0.1.4/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.47.0 h1:R1XyaNpoW4Et9yly+I2EeX7pBza/w+pmYee/0HJDyKk= -modernc.org/sqlite v1.47.0/go.mod h1:hWjRO6Tj/5Ik8ieqxQybiEOUXy0NJFNp2tpvVpKlvig= +modernc.org/sqlite v1.48.0 h1:ElZyLop3Q2mHYk5IFPPXADejZrlHu7APbpB0sF78bq4= +modernc.org/sqlite v1.48.0/go.mod h1:hWjRO6Tj/5Ik8ieqxQybiEOUXy0NJFNp2tpvVpKlvig= 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= diff --git a/pkg/agent/context_seahorse.go b/pkg/agent/context_seahorse.go index 104a84a78..327c6162a 100644 --- a/pkg/agent/context_seahorse.go +++ b/pkg/agent/context_seahorse.go @@ -1,3 +1,5 @@ +//go:build !mipsle && !netbsd && !(freebsd && arm) + package agent import ( diff --git a/pkg/agent/context_seahorse_unsupported.go b/pkg/agent/context_seahorse_unsupported.go new file mode 100644 index 000000000..7528f79bc --- /dev/null +++ b/pkg/agent/context_seahorse_unsupported.go @@ -0,0 +1,20 @@ +//go:build mipsle || netbsd || (freebsd && arm) + +package agent + +import ( + "encoding/json" + "fmt" +) + +// newSeahorseContextManager is unavailable on platforms where modernc sqlite/libc +// currently has no stable build path for this project. +func newSeahorseContextManager(_ json.RawMessage, _ *AgentLoop) (ContextManager, error) { + return nil, fmt.Errorf("seahorse context manager is unavailable on this platform") +} + +func init() { + if err := RegisterContextManager("seahorse", newSeahorseContextManager); err != nil { + panic(fmt.Sprintf("register seahorse context manager: %v", err)) + } +} diff --git a/pkg/agent/hook_process.go b/pkg/agent/hook_process.go index e5632913d..ace95f44d 100644 --- a/pkg/agent/hook_process.go +++ b/pkg/agent/hook_process.go @@ -12,7 +12,9 @@ import ( "sync/atomic" "time" + "github.com/sipeed/picoclaw/pkg/isolation" "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/tools" ) const ( @@ -90,7 +92,8 @@ type processHookAfterLLMResponse struct { type processHookBeforeToolResponse struct { processHookDecisionResponse - Call *ToolCallHookRequest `json:"call,omitempty"` + Call *ToolCallHookRequest `json:"call,omitempty"` + Result *tools.ToolResult `json:"result,omitempty"` // Result returned directly by hook (for respond action) } type processHookAfterToolResponse struct { @@ -120,7 +123,9 @@ func NewProcessHook(ctx context.Context, name string, opts ProcessHookOptions) ( if err != nil { return nil, fmt.Errorf("create process hook stderr: %w", err) } - if err := cmd.Start(); err != nil { + // Route hook subprocess startup through the shared isolation entry point so + // process hooks inherit the same isolation behavior as other child processes. + if err := isolation.Start(cmd); err != nil { return nil, fmt.Errorf("start process hook: %w", err) } @@ -241,6 +246,10 @@ func (ph *ProcessHook) BeforeTool( if resp.Call == nil { resp.Call = call } + // If hook returned a Result, carry it in ToolCallHookRequest + if resp.Result != nil { + resp.Call.HookResult = resp.Result + } return resp.Call, HookDecision{Action: resp.Action, Reason: resp.Reason}, nil } diff --git a/pkg/agent/hook_process_test.go b/pkg/agent/hook_process_test.go index 50f89811f..9e95d105e 100644 --- a/pkg/agent/hook_process_test.go +++ b/pkg/agent/hook_process_test.go @@ -7,10 +7,13 @@ import ( "fmt" "os" "path/filepath" + "runtime" "strings" "testing" "time" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/isolation" "github.com/sipeed/picoclaw/pkg/providers" ) @@ -178,6 +181,76 @@ func TestAgentLoop_MountProcessHook_ApprovalDeny(t *testing.T) { } } +func TestAgentLoop_MountProcessHook_IsolationSupportsRelativeDirAndCommand(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("linux-only isolation path handling") + } + + provider := &llmHookTestProvider{} + al, agent, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + root := t.TempDir() + t.Setenv(config.EnvHome, filepath.Join(root, "picoclaw-home")) + binDir := filepath.Join(root, "bin") + hookDir := filepath.Join(root, "hooks") + if err := os.MkdirAll(binDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.MkdirAll(hookDir, 0o755); err != nil { + t.Fatal(err) + } + writeFakeBwrap(t, filepath.Join(binDir, "bwrap")) + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + linkTestBinary(t, os.Args[0], filepath.Join(hookDir, "hook-helper")) + + cfg := config.DefaultConfig() + cfg.Isolation.Enabled = true + isolation.Configure(cfg) + t.Cleanup(func() { isolation.Configure(config.DefaultConfig()) }) + + cwd, err := os.Getwd() + if err != nil { + t.Fatal(err) + } + relHookDir, err := filepath.Rel(cwd, hookDir) + if err != nil { + t.Fatal(err) + } + + mountErr := al.MountProcessHook(context.Background(), "ipc-relative", ProcessHookOptions{ + Command: []string{"./hook-helper", "-test.run=TestProcessHook_HelperProcess", "--"}, + Dir: relHookDir, + Env: processHookHelperEnv("rewrite", ""), + InterceptLLM: true, + }) + if mountErr != nil { + t.Fatalf("MountProcessHook failed with relative dir/command under isolation: %v", mountErr) + } + + resp, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "session-relative", + Channel: "cli", + ChatID: "direct", + UserMessage: "hello", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + }) + if err != nil { + t.Fatalf("runAgentLoop failed: %v", err) + } + if resp != "provider content|ipc" { + t.Fatalf("expected process-hooked llm content, got %q", resp) + } + provider.mu.Lock() + lastModel := provider.lastModel + provider.mu.Unlock() + if lastModel != "process-model" { + t.Fatalf("expected process model, got %q", lastModel) + } +} + func processHookHelperCommand() []string { return []string{os.Args[0], "-test.run=TestProcessHook_HelperProcess", "--"} } @@ -193,6 +266,59 @@ func processHookHelperEnv(mode, eventLog string) []string { return env } +func writeFakeBwrap(t *testing.T, path string) { + t.Helper() + script := `#!/bin/sh +set -eu +workdir= +while [ "$#" -gt 0 ]; do + case "$1" in + --) + shift + break + ;; + --chdir) + workdir="$2" + shift 2 + ;; + --bind|--ro-bind) + shift 3 + ;; + --proc|--dev) + shift 2 + ;; + --die-with-parent|--unshare-ipc) + shift + ;; + *) + shift + ;; + esac +done +if [ -n "$workdir" ]; then + cd "$workdir" +fi +exec "$@" +` + if err := os.WriteFile(path, []byte(script), 0o755); err != nil { + t.Fatalf("write fake bwrap: %v", err) + } +} + +func linkTestBinary(t *testing.T, source, target string) { + t.Helper() + if err := os.Symlink(source, target); err == nil { + return + } + data, err := os.ReadFile(source) + if err != nil { + t.Fatalf("read test binary: %v", err) + } + if err := os.WriteFile(target, data, 0o755); err != nil { + t.Fatalf("create hook helper binary: %v", err) + } +} + func waitForFileContains(t *testing.T, path, substring string) { t.Helper() diff --git a/pkg/agent/hooks.go b/pkg/agent/hooks.go index c1ef58ffd..c23961dc6 100644 --- a/pkg/agent/hooks.go +++ b/pkg/agent/hooks.go @@ -25,6 +25,7 @@ type HookAction string const ( HookActionContinue HookAction = "continue" HookActionModify HookAction = "modify" + HookActionRespond HookAction = "respond" // Return result directly, skip tool execution. SECURITY: This bypasses ApproveTool checks, allowing hooks to return results for any tool (including sensitive ones like bash) without approval. Use with caution. HookActionDenyTool HookAction = "deny_tool" HookActionAbortTurn HookAction = "abort_turn" HookActionHardAbort HookAction = "hard_abort" @@ -127,11 +128,12 @@ func (r *LLMHookResponse) Clone() *LLMHookResponse { } type ToolCallHookRequest struct { - Meta EventMeta `json:"meta"` - Tool string `json:"tool"` - Arguments map[string]any `json:"arguments,omitempty"` - Channel string `json:"channel,omitempty"` - ChatID string `json:"chat_id,omitempty"` + Meta EventMeta `json:"meta"` + Tool string `json:"tool"` + Arguments map[string]any `json:"arguments,omitempty"` + Channel string `json:"channel,omitempty"` + ChatID string `json:"chat_id,omitempty"` + HookResult *tools.ToolResult `json:"hook_result,omitempty"` // Result returned directly by hook (for respond action). Media is supported - see Media handling section in docs. } func (r *ToolCallHookRequest) Clone() *ToolCallHookRequest { @@ -140,6 +142,7 @@ func (r *ToolCallHookRequest) Clone() *ToolCallHookRequest { } cloned := *r cloned.Arguments = cloneStringAnyMap(r.Arguments) + cloned.HookResult = cloneToolResult(r.HookResult) return &cloned } @@ -382,6 +385,10 @@ func (hm *HookManager) BeforeTool( if next != nil { current = next } + case HookActionRespond: + // Hook returns result directly, skip tool execution + // Carry HookResult in ToolCallHookRequest and return + return next, decision case HookActionDenyTool, HookActionAbortTurn, HookActionHardAbort: return current, decision default: @@ -793,6 +800,13 @@ func cloneToolResult(result *tools.ToolResult) *tools.ToolResult { if len(result.Media) > 0 { cloned.Media = append([]string(nil), result.Media...) } + if len(result.ArtifactTags) > 0 { + cloned.ArtifactTags = append([]string(nil), result.ArtifactTags...) + } + if len(result.Messages) > 0 { + cloned.Messages = make([]providers.Message, len(result.Messages)) + copy(cloned.Messages, result.Messages) + } return &cloned } diff --git a/pkg/agent/hooks_test.go b/pkg/agent/hooks_test.go index 49e1b1784..9049a5c72 100644 --- a/pkg/agent/hooks_test.go +++ b/pkg/agent/hooks_test.go @@ -2,6 +2,7 @@ package agent import ( "context" + "errors" "os" "sync" "testing" @@ -10,6 +11,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/tools" ) @@ -343,3 +345,517 @@ func TestAgentLoop_Hooks_ToolApproverCanDeny(t *testing.T) { t.Fatalf("expected skipped reason %q, got %q", expected, payload.Reason) } } + +// respondHook is a test hook for testing HookActionRespond functionality +type respondHook struct { + respondTools map[string]bool // tool names to respond to +} + +func (h *respondHook) BeforeTool( + ctx context.Context, + call *ToolCallHookRequest, +) (*ToolCallHookRequest, HookDecision, error) { + if h.respondTools[call.Tool] { + next := call.Clone() + next.HookResult = &tools.ToolResult{ + ForLLM: "hook-responded: " + call.Tool, + ForUser: "", + Silent: false, + IsError: false, + } + return next, HookDecision{Action: HookActionRespond}, nil + } + return call, HookDecision{Action: HookActionContinue}, nil +} + +func (h *respondHook) AfterTool( + ctx context.Context, + result *ToolResultHookResponse, +) (*ToolResultHookResponse, HookDecision, error) { + // Should not be called since respond skips tool execution + return result, HookDecision{Action: HookActionContinue}, nil +} + +func TestAgentLoop_Hooks_ToolRespondAction(t *testing.T) { + provider := &toolHookProvider{} + al, agent, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + al.RegisterTool(&echoTextTool{}) + if err := al.MountHook(NamedHook("respond-hook", &respondHook{ + respondTools: map[string]bool{"echo_text": true}, + })); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + sub := al.SubscribeEvents(16) + defer al.UnsubscribeEvents(sub.ID) + + resp, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "session-1", + Channel: "cli", + ChatID: "direct", + UserMessage: "run tool", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + }) + if err != nil { + t.Fatalf("runAgentLoop failed: %v", err) + } + + // Verify response comes from hook, not tool + expected := "hook-responded: echo_text" + if resp != expected { + t.Fatalf("expected %q, got %q", expected, resp) + } + + // Verify event stream has ToolExecEnd, not actual tool execution + events := collectEventStream(sub.C) + endEvt, ok := findEvent(events, EventKindToolExecEnd) + if !ok { + t.Fatal("expected tool exec end event") + } + payload, ok := endEvt.Payload.(ToolExecEndPayload) + if !ok { + t.Fatalf("expected ToolExecEndPayload, got %T", endEvt.Payload) + } + if payload.Tool != "echo_text" { + t.Fatalf("expected tool echo_text, got %q", payload.Tool) + } + if payload.ForLLMLen != len(expected) { + t.Fatalf("expected ForLLMLen %d, got %d", len(expected), payload.ForLLMLen) + } +} + +// denyToolHook tests HookActionDenyTool functionality +type denyToolHook struct { + denyTools map[string]bool +} + +func (h *denyToolHook) BeforeTool( + ctx context.Context, + call *ToolCallHookRequest, +) (*ToolCallHookRequest, HookDecision, error) { + if h.denyTools[call.Tool] { + return call, HookDecision{Action: HookActionDenyTool, Reason: "tool denied by hook"}, nil + } + return call, HookDecision{Action: HookActionContinue}, nil +} + +func (h *denyToolHook) AfterTool( + ctx context.Context, + result *ToolResultHookResponse, +) (*ToolResultHookResponse, HookDecision, error) { + return result, HookDecision{Action: HookActionContinue}, nil +} + +func TestAgentLoop_Hooks_ToolDenyAction(t *testing.T) { + provider := &toolHookProvider{} + al, agent, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + al.RegisterTool(&echoTextTool{}) + if err := al.MountHook(NamedHook("deny-hook", &denyToolHook{ + denyTools: map[string]bool{"echo_text": true}, + })); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + resp, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "session-1", + Channel: "cli", + ChatID: "direct", + UserMessage: "run tool", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + }) + if err != nil { + t.Fatalf("runAgentLoop failed: %v", err) + } + + expected := "Tool execution denied by hook: tool denied by hook" + if resp != expected { + t.Fatalf("expected %q, got %q", expected, resp) + } +} + +func TestHookManager_BeforeTool_RespondAction(t *testing.T) { + hm := NewHookManager(nil) + defer hm.Close() + + hook := &respondHook{ + respondTools: map[string]bool{"test_tool": true}, + } + if err := hm.Mount(NamedHook("respond-test", hook)); err != nil { + t.Fatalf("mount hook: %v", err) + } + + req := &ToolCallHookRequest{ + Tool: "test_tool", + Arguments: map[string]any{"arg": "value"}, + } + result, decision := hm.BeforeTool(context.Background(), req) + + if decision.Action != HookActionRespond { + t.Fatalf("expected action %q, got %q", HookActionRespond, decision.Action) + } + + if result.HookResult == nil { + t.Fatal("expected HookResult to be set") + } + if result.HookResult.ForLLM != "hook-responded: test_tool" { + t.Fatalf("unexpected HookResult.ForLLM: %q", result.HookResult.ForLLM) + } +} + +type respondWithMediaHook struct { + respondTools map[string]bool + media []string + responseHandled bool + forLLM string +} + +func (h *respondWithMediaHook) BeforeTool( + ctx context.Context, + call *ToolCallHookRequest, +) (*ToolCallHookRequest, HookDecision, error) { + if h.respondTools[call.Tool] { + next := call.Clone() + next.HookResult = &tools.ToolResult{ + ForLLM: h.forLLM, + ForUser: "media result", + Media: h.media, + ResponseHandled: h.responseHandled, + Silent: false, + IsError: false, + } + return next, HookDecision{Action: HookActionRespond}, nil + } + return call, HookDecision{Action: HookActionContinue}, nil +} + +func (h *respondWithMediaHook) AfterTool( + ctx context.Context, + result *ToolResultHookResponse, +) (*ToolResultHookResponse, HookDecision, error) { + return result, HookDecision{Action: HookActionContinue}, nil +} + +type errorMediaChannel struct { + fakeChannel + sendErr error +} + +func (f *errorMediaChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) ([]string, error) { + return nil, f.sendErr +} + +func TestAgentLoop_HookRespond_MediaError(t *testing.T) { + provider := &multiToolProvider{ + toolCalls: []providers.ToolCall{ + {ID: "call-1", Name: "media_tool", Arguments: map[string]any{}}, + }, + finalContent: "done", + } + al, agent, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + hook := &respondWithMediaHook{ + respondTools: map[string]bool{"media_tool": true}, + media: []string{"media://test/image.png"}, + responseHandled: true, + forLLM: "media sent successfully", + } + if err := al.MountHook(NamedHook("media-hook", hook)); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + al.channelManager = newStartedTestChannelManager(t, al.bus, al.mediaStore, "discord", &errorMediaChannel{ + sendErr: errors.New("channel unavailable"), + }) + + sub := al.SubscribeEvents(16) + defer al.UnsubscribeEvents(sub.ID) + + _, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "session-media-err", + Channel: "discord", + ChatID: "chat1", + UserMessage: "send media", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + }) + if err != nil { + t.Fatalf("runAgentLoop failed: %v", err) + } + + events := collectEventStream(sub.C) + endEvt, ok := findEvent(events, EventKindToolExecEnd) + if !ok { + t.Fatal("expected ToolExecEnd event") + } + payload, ok := endEvt.Payload.(ToolExecEndPayload) + if !ok { + t.Fatalf("expected ToolExecEndPayload, got %T", endEvt.Payload) + } + + if !payload.IsError { + t.Fatal("expected IsError=true when SendMedia fails") + } + + if payload.ForLLMLen < 30 { + t.Fatalf("expected ForLLM to contain error message, got ForLLMLen=%d", payload.ForLLMLen) + } +} + +func TestAgentLoop_HookRespond_BusFallback(t *testing.T) { + provider := &multiToolProvider{ + toolCalls: []providers.ToolCall{ + {ID: "call-1", Name: "media_tool", Arguments: map[string]any{}}, + }, + finalContent: "done", + } + al, agent, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + hook := &respondWithMediaHook{ + respondTools: map[string]bool{"media_tool": true}, + media: []string{"media://test/image.png"}, + responseHandled: true, + forLLM: "media queued", + } + if err := al.MountHook(NamedHook("media-hook", hook)); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + sub := al.SubscribeEvents(16) + defer al.UnsubscribeEvents(sub.ID) + + resp, err := al.runAgentLoop(context.Background(), agent, processOptions{ + SessionKey: "session-bus-fallback", + Channel: "cli", + ChatID: "chat1", + UserMessage: "send media", + DefaultResponse: defaultResponse, + EnableSummary: false, + SendResponse: false, + }) + if err != nil { + t.Fatalf("runAgentLoop failed: %v", err) + } + + events := collectEventStream(sub.C) + endEvt, ok := findEvent(events, EventKindToolExecEnd) + if !ok { + t.Fatal("expected ToolExecEnd event") + } + payload, ok := endEvt.Payload.(ToolExecEndPayload) + if !ok { + t.Fatalf("expected ToolExecEndPayload, got %T", endEvt.Payload) + } + + if payload.IsError { + t.Fatal("expected IsError=false for bus fallback (media queued, not delivered)") + } + + if resp != "done" { + t.Fatalf("expected response 'done', got %q", resp) + } +} + +type multiToolProvider struct { + mu sync.Mutex + callCount int + toolCalls []providers.ToolCall + finalContent string +} + +func (p *multiToolProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + p.mu.Lock() + defer p.mu.Unlock() + + p.callCount++ + if p.callCount == 1 && len(p.toolCalls) > 0 { + return &providers.LLMResponse{ + ToolCalls: p.toolCalls, + }, nil + } + + return &providers.LLMResponse{ + Content: p.finalContent, + }, nil +} + +func (p *multiToolProvider) GetDefaultModel() string { + return "multi-tool-provider" +} + +func TestAgentLoop_HookRespond_InterruptSkipsRemaining(t *testing.T) { + provider := &multiToolProvider{ + toolCalls: []providers.ToolCall{ + {ID: "call-1", Name: "tool_one", Arguments: map[string]any{}}, + {ID: "call-2", Name: "tool_two", Arguments: map[string]any{}}, + {ID: "call-3", Name: "tool_three", Arguments: map[string]any{}}, + }, + finalContent: "done", + } + al, _, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + tool1ExecCh := make(chan struct{}, 1) + al.RegisterTool(&slowTool{name: "tool_two", duration: 100 * time.Millisecond, execCh: tool1ExecCh}) + al.RegisterTool(&slowTool{name: "tool_three", duration: 100 * time.Millisecond}) + + hook := &respondHook{ + respondTools: map[string]bool{"tool_one": true}, + } + if err := al.MountHook(NamedHook("respond-hook", hook)); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + sub := al.SubscribeEvents(32) + defer al.UnsubscribeEvents(sub.ID) + + sessionKey := routing.BuildAgentMainSessionKey(routing.DefaultAgentID) + + type result struct { + resp string + err error + } + resultCh := make(chan result, 1) + go func() { + resp, err := al.ProcessDirectWithChannel( + context.Background(), + "run tools", + sessionKey, + "cli", + "chat1", + ) + resultCh <- result{resp: resp, err: err} + }() + + time.Sleep(50 * time.Millisecond) + + if err := al.InterruptGraceful("stop now"); err != nil { + t.Fatalf("InterruptGraceful failed: %v", err) + } + + select { + case r := <-resultCh: + if r.err != nil { + t.Fatalf("unexpected error: %v", r.err) + } + case <-time.After(3 * time.Second): + t.Fatal("timeout waiting for result") + } + + events := collectEventStream(sub.C) + + skippedEvts := filterEvents(events, EventKindToolExecSkipped) + if len(skippedEvts) < 1 { + t.Fatal("expected at least one ToolExecSkipped event after interrupt") + } + + for _, evt := range skippedEvts { + payload, ok := evt.Payload.(ToolExecSkippedPayload) + if !ok { + t.Fatalf("expected ToolExecSkippedPayload, got %T", evt.Payload) + } + if payload.Reason != "graceful interrupt requested" { + t.Fatalf("expected skip reason 'graceful interrupt requested', got %q", payload.Reason) + } + } +} + +func TestAgentLoop_HookRespond_SteeringSkipsRemaining(t *testing.T) { + provider := &multiToolProvider{ + toolCalls: []providers.ToolCall{ + {ID: "call-1", Name: "tool_one", Arguments: map[string]any{}}, + {ID: "call-2", Name: "tool_two", Arguments: map[string]any{}}, + {ID: "call-3", Name: "tool_three", Arguments: map[string]any{}}, + }, + finalContent: "done", + } + al, _, cleanup := newHookTestLoop(t, provider) + defer cleanup() + + al.RegisterTool(&slowTool{name: "tool_two", duration: 100 * time.Millisecond}) + al.RegisterTool(&slowTool{name: "tool_three", duration: 100 * time.Millisecond}) + + hook := &respondHook{ + respondTools: map[string]bool{"tool_one": true}, + } + if err := al.MountHook(NamedHook("respond-hook", hook)); err != nil { + t.Fatalf("MountHook failed: %v", err) + } + + sub := al.SubscribeEvents(32) + defer al.UnsubscribeEvents(sub.ID) + + sessionKey := routing.BuildAgentMainSessionKey(routing.DefaultAgentID) + + type result struct { + resp string + err error + } + resultCh := make(chan result, 1) + go func() { + resp, err := al.ProcessDirectWithChannel( + context.Background(), + "run tools", + sessionKey, + "cli", + "chat1", + ) + resultCh <- result{resp: resp, err: err} + }() + + time.Sleep(50 * time.Millisecond) + + al.Steer(providers.Message{Role: "user", Content: "change direction"}) + + select { + case r := <-resultCh: + if r.err != nil { + t.Fatalf("unexpected error: %v", r.err) + } + case <-time.After(3 * time.Second): + t.Fatal("timeout waiting for result") + } + + events := collectEventStream(sub.C) + + skippedEvts := filterEvents(events, EventKindToolExecSkipped) + if len(skippedEvts) < 1 { + t.Fatal("expected at least one ToolExecSkipped event after steering") + } + + for _, evt := range skippedEvts { + payload, ok := evt.Payload.(ToolExecSkippedPayload) + if !ok { + t.Fatalf("expected ToolExecSkippedPayload, got %T", evt.Payload) + } + if payload.Reason != "queued user steering message" { + t.Fatalf("expected skip reason 'queued user steering message', got %q", payload.Reason) + } + } +} + +func filterEvents(events []Event, kind EventKind) []Event { + var result []Event + for _, evt := range events { + if evt.Kind == kind { + result = append(result, evt) + } + } + return result +} diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index bacfa49c5..5bcb83087 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -9,6 +9,7 @@ import ( "strings" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/isolation" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/memory" @@ -51,6 +52,10 @@ type AgentInstance struct { // LightProvider is the concrete provider instance for the configured light model. // It is only used when routing selects the light tier for a turn. LightProvider providers.LLMProvider + // CandidateProviders maps "provider/model" keys to per-candidate LLMProvider + // instances. This allows each fallback model to use its own api_base and api_key + // from model_list, instead of inheriting the primary model's provider config. + CandidateProviders map[string]providers.LLMProvider } // NewAgentInstance creates an agent instance from config. @@ -60,6 +65,12 @@ func NewAgentInstance( cfg *config.Config, provider providers.LLMProvider, ) *AgentInstance { + if cfg != nil { + // Keep the subprocess isolation runtime aligned with the latest loaded config + // before any tools or providers start spawning child processes. + isolation.Configure(cfg) + } + workspace := resolveAgentWorkspace(agentCfg, defaults) os.MkdirAll(workspace, 0o755) @@ -175,6 +186,9 @@ func NewAgentInstance( // Resolve fallback candidates candidates := resolveModelCandidates(cfg, defaults.Provider, model, fallbacks) + candidateProviders := make(map[string]providers.LLMProvider) + populateCandidateProvidersFromNames(cfg, workspace, fallbacks, candidateProviders) + // Model routing setup: pre-resolve light model candidates at creation time // to avoid repeated model_list lookups on every incoming message. var router *routing.Router @@ -199,6 +213,7 @@ func NewAgentInstance( }) lightCandidates = resolved lightProvider = lp + populateCandidateProvidersFromNames(cfg, workspace, []string{rc.LightModel}, candidateProviders) } } } else { @@ -230,6 +245,43 @@ func NewAgentInstance( Router: router, LightCandidates: lightCandidates, LightProvider: lightProvider, + CandidateProviders: candidateProviders, + } +} + +// populateCandidateProvidersFromNames resolves each model name (alias or +// "provider/model") via resolvedModelConfig and creates a dedicated LLMProvider +// for it. This reuses the canonical config resolution path (GetModelConfig) so +// alias handling and load-balancing stay consistent with the rest of the codebase. +func populateCandidateProvidersFromNames( + cfg *config.Config, + workspace string, + names []string, + out map[string]providers.LLMProvider, +) { + if cfg == nil || len(names) == 0 { + return + } + for _, name := range names { + mc, err := resolvedModelConfig(cfg, strings.TrimSpace(name), workspace) + if err != nil { + logger.WarnCF("agent", + "fallback provider: no model_list entry found; will inherit primary provider credentials", + map[string]any{"name": name, "error": err.Error()}) + continue + } + protocol, modelID := providers.ExtractProtocol(strings.TrimSpace(mc.Model)) + key := providers.ModelKey(providers.NormalizeProvider(protocol), modelID) + if _, exists := out[key]; exists { + continue + } + p, _, err := providers.CreateProviderFromConfig(mc) + if err != nil { + logger.WarnCF("agent", "fallback provider: failed to create provider", + map[string]any{"model": mc.Model, "error": err.Error()}) + continue + } + out[key] = p } } diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index ba907e88b..8c71296ed 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -9,6 +9,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/media" + "github.com/sipeed/picoclaw/pkg/providers" ) func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) { @@ -300,6 +301,199 @@ func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) { } } +// TestPopulateCandidateProviders_NilCfgIsNoop verifies that passing a nil +// config does not panic and leaves the output map empty. +func TestPopulateCandidateProviders_NilCfgIsNoop(t *testing.T) { + out := map[string]providers.LLMProvider{} + populateCandidateProvidersFromNames(nil, t.TempDir(), []string{"gpt-4o"}, out) + if len(out) != 0 { + t.Fatalf("expected empty map, got %d entries", len(out)) + } +} + +// TestPopulateCandidateProviders_SkipsExistingKeys verifies that a key already +// present in the output map is not overwritten. +func TestPopulateCandidateProviders_SkipsExistingKeys(t *testing.T) { + existing := &mockProvider{} + key := providers.ModelKey("openai", "gpt-4o") + out := map[string]providers.LLMProvider{key: existing} + + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + {ModelName: "my-gpt", Model: "openai/gpt-4o", APIKeys: config.SimpleSecureStrings("test-key")}, + }, + } + populateCandidateProvidersFromNames(cfg, t.TempDir(), []string{"my-gpt"}, out) + + if out[key] != existing { + t.Fatal("existing provider entry was overwritten; expected it to be preserved") + } +} + +// TestPopulateCandidateProviders_ResolvesAlias verifies that a model_name +// alias (e.g. "my-gpt") is resolved via GetModelConfig and the provider +// is created using the underlying model's config. +func TestPopulateCandidateProviders_ResolvesAlias(t *testing.T) { + workspace := t.TempDir() + out := map[string]providers.LLMProvider{} + + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + {ModelName: "my-gpt", Model: "openai/gpt-4o", APIBase: "https://api.openai.com/v1", Workspace: workspace}, + }, + } + populateCandidateProvidersFromNames(cfg, workspace, []string{"my-gpt"}, out) + + key := providers.ModelKey("openai", "gpt-4o") + if out[key] == nil { + t.Fatalf("expected CandidateProviders[%q] to be populated for alias", key) + } +} + +// TestPopulateCandidateProviders_ResolvesProtocolPrefix verifies that a +// model_list entry using full "provider/model" notation (e.g. +// "gemini/gemma-3-27b-it") is matched correctly when referenced by model_name. +func TestPopulateCandidateProviders_ResolvesProtocolPrefix(t *testing.T) { + workspace := t.TempDir() + out := map[string]providers.LLMProvider{} + + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + { + ModelName: "gemma", + Model: "gemini/gemma-3-27b-it", + APIKeys: config.SimpleSecureStrings("gemini-test-key"), + Workspace: workspace, + }, + }, + } + populateCandidateProvidersFromNames(cfg, workspace, []string{"gemma"}, out) + + key := providers.ModelKey("gemini", "gemma-3-27b-it") + if out[key] == nil { + t.Fatalf("expected CandidateProviders[%q] to be populated for protocol-prefixed model", key) + } +} + +// TestPopulateCandidateProviders_EmptyNamesIsNoop verifies the early-exit +// path when the names slice is empty. +func TestPopulateCandidateProviders_EmptyNamesIsNoop(t *testing.T) { + out := map[string]providers.LLMProvider{} + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + {ModelName: "my-gpt", Model: "openai/gpt-4o", APIKeys: config.SimpleSecureStrings("key")}, + }, + } + populateCandidateProvidersFromNames(cfg, t.TempDir(), nil, out) + if len(out) != 0 { + t.Fatalf("expected empty map, got %d entries", len(out)) + } +} + +// TestPopulateCandidateProviders_EmptyModelListIsNoop verifies the early-exit +// path when model_list is empty — no provider can be created. +func TestPopulateCandidateProviders_EmptyModelListIsNoop(t *testing.T) { + out := map[string]providers.LLMProvider{} + cfg := &config.Config{} + populateCandidateProvidersFromNames(cfg, t.TempDir(), []string{"gpt-4o"}, out) + if len(out) != 0 { + t.Fatalf("expected empty map, got %d entries", len(out)) + } +} + +// TestPopulateCandidateProviders_UnmatchedNameIsSkipped verifies that a +// name with no matching model_list entry is skipped and does not +// cause a panic or leave a nil entry in the map. +func TestPopulateCandidateProviders_UnmatchedNameIsSkipped(t *testing.T) { + out := map[string]providers.LLMProvider{} + cfg := &config.Config{ + ModelList: []*config.ModelConfig{ + {ModelName: "my-gpt", Model: "openai/gpt-4o", APIKeys: config.SimpleSecureStrings("key")}, + }, + } + populateCandidateProvidersFromNames(cfg, t.TempDir(), []string{"nonexistent-model"}, out) + + if len(out) != 0 { + t.Fatalf("expected empty map for unmatched name, got %d entries", len(out)) + } +} + +// TestNewAgentInstance_CandidateProvidersPopulatedForCrossProviderFallbacks +// mirrors the exact scenario from bug #2140: primary model on OpenRouter with +// Gemini fallbacks. Each entry must get its own provider instance so that +// fallback requests go to the correct API endpoint, not the primary's. +func TestNewAgentInstance_CandidateProvidersPopulatedForCrossProviderFallbacks(t *testing.T) { + workspace := t.TempDir() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "mistral-small-3.1", + ModelFallbacks: []string{"gemma-3-27b", "gemini-images"}, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "mistral-small-3.1", + Model: "openrouter/mistralai/mistral-small-3.1-24b-instruct:free", + APIBase: "https://openrouter.ai/api/v1", + APIKeys: config.SimpleSecureStrings("sk-or-test"), + Workspace: workspace, + }, + { + ModelName: "gemma-3-27b", + Model: "gemini/gemma-3-27b-it", + APIKeys: config.SimpleSecureStrings("AIzaSy-test"), + Workspace: workspace, + }, + { + ModelName: "gemini-images", + Model: "gemini/gemini-2.5-flash-lite", + APIKeys: config.SimpleSecureStrings("AIzaSy-test"), + Workspace: workspace, + }, + }, + } + + primaryProvider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, primaryProvider) + + // Only fallback models need entries — the primary uses the injected provider directly. + wantKeys := []string{ + providers.ModelKey("gemini", "gemma-3-27b-it"), + providers.ModelKey("gemini", "gemini-2.5-flash-lite"), + } + + for _, key := range wantKeys { + p, ok := agent.CandidateProviders[key] + if !ok { + t.Errorf("CandidateProviders missing key %q", key) + continue + } + if p == nil { + t.Errorf("CandidateProviders[%q] is nil", key) + } + // Each fallback must use its own provider, not the injected primary. + if p == primaryProvider { + t.Errorf( + "CandidateProviders[%q] is the same instance as the primary provider; fallback would inherit primary credentials", + key, + ) + } + } + + if t.Failed() { + t.Logf("CandidateProviders keys present: %v", func() []string { + keys := make([]string, 0, len(agent.CandidateProviders)) + for k := range agent.CandidateProviders { + keys = append(keys, k) + } + return keys + }()) + } +} + func TestNewAgentInstance_ReadFileModeSelectsSchema(t *testing.T) { workspace := t.TempDir() diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 8bee6f56c..5172d7ec9 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -673,21 +673,21 @@ func (al *AgentLoop) PublishResponseIfNeeded(ctx context.Context, channel, chatI return } - alreadySent := false + alreadySentToSameChat := false defaultAgent := al.GetRegistry().GetDefaultAgent() if defaultAgent != nil { if tool, ok := defaultAgent.Tools.Get("message"); ok { if mt, ok := tool.(*tools.MessageTool); ok { - alreadySent = mt.HasSentInRound() + alreadySentToSameChat = mt.HasSentTo(channel, chatID) } } } - if alreadySent { + if alreadySentToSameChat { logger.DebugCF( "agent", - "Skipped outbound (message tool already sent)", - map[string]any{"channel": channel}, + "Skipped outbound (message tool already sent to same chat)", + map[string]any{"channel": channel, "chat_id": chatID}, ) return } @@ -2088,7 +2088,11 @@ turnLoop: providerCtx, activeCandidates, func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { - return activeProvider.Chat(ctx, messagesForCall, toolDefsForCall, model, llmOpts) + candidateProvider := activeProvider + if cp, ok := ts.agent.CandidateProviders[providers.ModelKey(provider, model)]; ok { + candidateProvider = cp + } + return candidateProvider.Chat(ctx, messagesForCall, toolDefsForCall, model, llmOpts) }, ) if fbErr != nil { @@ -2421,6 +2425,236 @@ turnLoop: toolName = toolReq.Tool toolArgs = toolReq.Arguments } + case HookActionRespond: + // Hook returns result directly, skip tool execution. + // SECURITY: This bypasses ApproveTool, allowing hooks to respond + // for any tool name without approval. This is intentional for + // plugin tools but means a before_tool hook can override even + // sensitive tools like bash. Hook configuration should be + // carefully reviewed to prevent unauthorized tool execution. + if toolReq != nil && toolReq.HookResult != nil { + hookResult := toolReq.HookResult + + argsJSON, _ := json.Marshal(toolArgs) + argsPreview := utils.Truncate(string(argsJSON), 200) + logger.InfoCF("agent", fmt.Sprintf("Tool call (hook respond): %s(%s)", toolName, argsPreview), + map[string]any{ + "agent_id": ts.agent.ID, + "tool": toolName, + "iteration": iteration, + }) + + // Emit ToolExecStart event (same as normal tool execution) + al.emitEvent( + EventKindToolExecStart, + ts.eventMeta("runTurn", "turn.tool.start"), + ToolExecStartPayload{ + Tool: toolName, + Arguments: cloneEventArguments(toolArgs), + }, + ) + + // Send tool feedback to chat channel if enabled (same as normal tool execution) + if al.cfg.Agents.Defaults.IsToolFeedbackEnabled() && + ts.channel != "" && + !ts.opts.SuppressToolFeedback { + argsJSON, _ := json.Marshal(toolArgs) + feedbackPreview := utils.Truncate( + string(argsJSON), + al.cfg.Agents.Defaults.GetToolFeedbackMaxArgsLength(), + ) + feedbackMsg := fmt.Sprintf("\U0001f527 `%s`\n```\n%s\n```", toolName, feedbackPreview) + fbCtx, fbCancel := context.WithTimeout(turnCtx, 3*time.Second) + _ = al.bus.PublishOutbound(fbCtx, bus.OutboundMessage{ + Channel: ts.channel, + ChatID: ts.chatID, + Content: feedbackMsg, + }) + fbCancel() + } + + toolDuration := time.Duration(0) // Hook execution time unknown + + // Send ForUser content to user + // For ResponseHandled results, send regardless of SendResponse setting, + // same as normal tool execution path. + shouldSendForUser := !hookResult.Silent && hookResult.ForUser != "" && + (ts.opts.SendResponse || hookResult.ResponseHandled) + if shouldSendForUser { + al.bus.PublishOutbound(ctx, bus.OutboundMessage{ + Channel: ts.channel, + ChatID: ts.chatID, + Content: hookResult.ForUser, + Metadata: map[string]string{ + "is_tool_call": "true", + }, + }) + } + + // Handle media from hook result (same as normal tool execution) + if len(hookResult.Media) > 0 && hookResult.ResponseHandled { + parts := make([]bus.MediaPart, 0, len(hookResult.Media)) + for _, ref := range hookResult.Media { + part := bus.MediaPart{Ref: ref} + if al.mediaStore != nil { + if _, meta, err := al.mediaStore.ResolveWithMeta(ref); err == nil { + part.Filename = meta.Filename + part.ContentType = meta.ContentType + part.Type = inferMediaType(meta.Filename, meta.ContentType) + } + } + parts = append(parts, part) + } + outboundMedia := bus.OutboundMediaMessage{ + Channel: ts.channel, + ChatID: ts.chatID, + Parts: parts, + } + if al.channelManager != nil && ts.channel != "" && !constants.IsInternalChannel(ts.channel) { + if err := al.channelManager.SendMedia(ctx, outboundMedia); err != nil { + logger.WarnCF("agent", "Failed to deliver hook media", + map[string]any{ + "agent_id": ts.agent.ID, + "tool": toolName, + "channel": ts.channel, + "chat_id": ts.chatID, + "error": err.Error(), + }) + // Same as normal tool execution: notify LLM about delivery failure + hookResult.IsError = true + hookResult.ForLLM = fmt.Sprintf("failed to deliver attachment: %v", err) + } + } else if al.bus != nil { + al.bus.PublishOutboundMedia(ctx, outboundMedia) + // Same as normal tool execution: bus only queues, media not yet delivered + hookResult.ResponseHandled = false + } + } + + // Track response handling status (same as normal tool execution) + if !hookResult.ResponseHandled { + allResponsesHandled = false + } + + // Build tool message + contentForLLM := hookResult.ContentForLLM() + if al.cfg.Tools.IsFilterSensitiveDataEnabled() { + contentForLLM = al.cfg.FilterSensitiveData(contentForLLM) + } + + toolResultMsg := providers.Message{ + Role: "tool", + Content: contentForLLM, + ToolCallID: tc.ID, + } + + // Handle media for LLM vision (same as normal tool execution) + if len(hookResult.Media) > 0 && !hookResult.ResponseHandled { + hookResult.ArtifactTags = buildArtifactTags(al.mediaStore, hookResult.Media) + // Recalculate contentForLLM after adding ArtifactTags + contentForLLM = hookResult.ContentForLLM() + if al.cfg.Tools.IsFilterSensitiveDataEnabled() { + contentForLLM = al.cfg.FilterSensitiveData(contentForLLM) + } + toolResultMsg.Content = contentForLLM + toolResultMsg.Media = append(toolResultMsg.Media, hookResult.Media...) + } + + // Emit ToolExecEnd event (after filtering, same as normal tool execution) + al.emitEvent( + EventKindToolExecEnd, + ts.eventMeta("runTurn", "turn.tool.end"), + ToolExecEndPayload{ + Tool: toolName, + Duration: toolDuration, + ForLLMLen: len(contentForLLM), + ForUserLen: len(hookResult.ForUser), + IsError: hookResult.IsError, + Async: hookResult.Async, + }, + ) + + messages = append(messages, toolResultMsg) + if !ts.opts.NoHistory { + ts.agent.Sessions.AddFullMessage(ts.sessionKey, toolResultMsg) + ts.recordPersistedMessage(toolResultMsg) + ts.ingestMessage(turnCtx, al, toolResultMsg) + } + + // Same as normal tool execution: check for steering/interrupt/SubTurn after each tool + if steerMsgs := al.dequeueSteeringMessagesForScope(ts.sessionKey); len(steerMsgs) > 0 { + pendingMessages = append(pendingMessages, steerMsgs...) + } + + skipReason := "" + skipMessage := "" + if len(pendingMessages) > 0 { + skipReason = "queued user steering message" + skipMessage = "Skipped due to queued user message." + } else if gracefulPending, _ := ts.gracefulInterruptRequested(); gracefulPending { + skipReason = "graceful interrupt requested" + skipMessage = "Skipped due to graceful interrupt." + } + + if skipReason != "" { + remaining := len(normalizedToolCalls) - i - 1 + if remaining > 0 { + logger.InfoCF("agent", "Turn checkpoint: skipping remaining tools after hook respond", + map[string]any{ + "agent_id": ts.agent.ID, + "completed": i + 1, + "skipped": remaining, + "reason": skipReason, + }) + for j := i + 1; j < len(normalizedToolCalls); j++ { + skippedTC := normalizedToolCalls[j] + al.emitEvent( + EventKindToolExecSkipped, + ts.eventMeta("runTurn", "turn.tool.skipped"), + ToolExecSkippedPayload{ + Tool: skippedTC.Name, + Reason: skipReason, + }, + ) + skippedMsg := providers.Message{ + Role: "tool", + Content: skipMessage, + ToolCallID: skippedTC.ID, + } + messages = append(messages, skippedMsg) + if !ts.opts.NoHistory { + ts.agent.Sessions.AddFullMessage(ts.sessionKey, skippedMsg) + ts.recordPersistedMessage(skippedMsg) + } + } + } + break + } + + // Also poll for any SubTurn results that arrived during tool execution. + if ts.pendingResults != nil { + select { + case result, ok := <-ts.pendingResults: + if ok && result != nil && result.ForLLM != "" { + content := al.cfg.FilterSensitiveData(result.ForLLM) + msg := providers.Message{Role: "user", Content: fmt.Sprintf("[SubTurn Result] %s", content)} + messages = append(messages, msg) + ts.agent.Sessions.AddFullMessage(ts.sessionKey, msg) + } + default: + // No results available + } + } + + continue + } + // If no HookResult, fall back to continue with warning + logger.WarnCF("agent", "Hook returned respond action but no HookResult provided", + map[string]any{ + "agent_id": ts.agent.ID, + "tool": toolName, + "action": "respond", + }) case HookActionDenyTool: allResponsesHandled = false denyContent := hookDeniedToolContent("Tool execution denied by hook", decision.Reason) diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 86a5c0fc2..ea91910d7 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -1840,6 +1840,164 @@ func TestProcessMessage_ModelRoutingUsesLightProvider(t *testing.T) { } } +// TestProcessMessage_FallbackUsesPerCandidateProvider is the loop-level test for +// bug #2140. It verifies that when the primary model returns a rate-limit error +// the fallback closure routes the retry to the fallback model's own provider +// (its own api_base), not back to the primary provider's endpoint. +func TestProcessMessage_FallbackUsesPerCandidateProvider(t *testing.T) { + workspace := t.TempDir() + + primaryCalls := 0 + primaryServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + primaryCalls++ + // Return 429 so FallbackChain classifies this as retriable and moves on. + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{ + "message": "rate limit exceeded", + "type": "rate_limit_error", + }, + }) + })) + defer primaryServer.Close() + + fallbackCalls := 0 + fallbackServer := newStrictChatCompletionTestServer( + t, "fallback", "gemma-3-27b-it", "fallback reply", &fallbackCalls, + ) + defer fallbackServer.Close() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "mistral-primary", + ModelFallbacks: []string{"gemma-fallback"}, + MaxTokens: 4096, + MaxToolIterations: 3, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "mistral-primary", + Model: "openrouter/mistralai/mistral-small-3.1", + APIBase: primaryServer.URL, + APIKeys: config.SimpleSecureStrings("primary-key"), + Workspace: workspace, + }, + { + ModelName: "gemma-fallback", + Model: "gemini/gemma-3-27b-it", + APIBase: fallbackServer.URL, + APIKeys: config.SimpleSecureStrings("fallback-key"), + Workspace: workspace, + }, + }, + } + + provider, _, err := providers.CreateProvider(cfg) + if err != nil { + t.Fatalf("CreateProvider() error = %v", err) + } + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) + helper := testHelper{al: al} + + resp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "hi", + Peer: bus.Peer{Kind: "direct", ID: "user1"}, + }) + + if resp != "fallback reply" { + t.Fatalf("response = %q, want %q (fallback provider)", resp, "fallback reply") + } + if primaryCalls == 0 { + t.Fatal("primary server was never called; expected at least one attempt") + } + if fallbackCalls != 1 { + t.Fatalf("fallback server calls = %d, want 1", fallbackCalls) + } +} + +// TestProcessMessage_FallbackUsesActiveProviderWhenCandidateNotRegistered verifies +// that when a candidate has no model_list entry it is absent from CandidateProviders +// and the fallback closure falls back to activeProvider instead of panicking. +func TestProcessMessage_FallbackUsesActiveProviderWhenCandidateNotRegistered(t *testing.T) { + workspace := t.TempDir() + + // Primary server: returns 429 on first call, succeeds on second. + // Both the primary and the unregistered fallback share this server + // (same api_base) so activeProvider routes both calls here. + callCount := 0 + primaryServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + callCount++ + w.Header().Set("Content-Type", "application/json") + if callCount == 1 { + w.WriteHeader(http.StatusTooManyRequests) + _ = json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{"message": "rate limit", "type": "rate_limit_error"}, + }) + return + } + // Second call (fallback via activeProvider) succeeds. + _ = json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "active provider reply"}, "finish_reason": "stop"}, + }, + }) + })) + defer primaryServer.Close() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: workspace, + ModelName: "primary-model", + MaxTokens: 4096, + MaxToolIterations: 3, + // No model_list entry for this alias — absent from CandidateProviders. + ModelFallbacks: []string{"openrouter/fallback-model"}, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "primary-model", + Model: "openrouter/primary-model", + APIBase: primaryServer.URL, + APIKeys: config.SimpleSecureStrings("primary-key"), + Workspace: workspace, + }, + }, + } + + provider, _, err := providers.CreateProvider(cfg) + if err != nil { + t.Fatalf("CreateProvider() error = %v", err) + } + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) + + helper := testHelper{al: al} + resp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "hi", + Peer: bus.Peer{Kind: "direct", ID: "user1"}, + }) + + if resp != "active provider reply" { + t.Fatalf("response = %q, want %q", resp, "active provider reply") + } + if callCount < 2 { + t.Fatalf("primary server calls = %d, want >= 2 (one 429 + one success via activeProvider)", callCount) + } +} + // TestToolResult_SilentToolDoesNotSendUserMessage verifies silent tools don't trigger outbound func TestToolResult_SilentToolDoesNotSendUserMessage(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") diff --git a/pkg/channels/feishu/feishu_64.go b/pkg/channels/feishu/feishu_64.go index b0b231d09..c12827729 100644 --- a/pkg/channels/feishu/feishu_64.go +++ b/pkg/channels/feishu/feishu_64.go @@ -14,6 +14,7 @@ import ( "strings" "sync" "sync/atomic" + "time" lark "github.com/larksuite/oapi-sdk-go/v3" larkcore "github.com/larksuite/oapi-sdk-go/v3/core" @@ -42,12 +43,18 @@ type FeishuChannel struct { wsClient *larkws.Client tokenCache *tokenCache // custom cache that supports invalidation - botOpenID atomic.Value // stores string; populated lazily for @mention detection + botOpenID atomic.Value // stores string; populated lazily for @mention detection + messageCache sync.Map // caches fetched messages (messageID -> *larkim.Message) mu sync.Mutex cancel context.CancelFunc } +type cachedMessage struct { + msg *larkim.Message + expiry time.Time +} + func NewFeishuChannel(cfg config.FeishuConfig, bus *bus.MessageBus) (*FeishuChannel, error) { base := channels.NewBaseChannel("feishu", cfg, bus, cfg.AllowFrom, channels.WithGroupTrigger(cfg.GroupTrigger), @@ -436,24 +443,8 @@ func (c *FeishuChannel) handleMessageReceive(ctx context.Context, event *larkim. // Append media tags to content (like Telegram does) content = appendMediaTags(content, messageType, mediaRefs) - if content == "" { - content = "[empty message]" - } - - metadata := map[string]string{} - if messageID != "" { - metadata["message_id"] = messageID - } - if messageType != "" { - metadata["message_type"] = messageType - } chatType := stringValue(message.ChatType) - if chatType != "" { - metadata["chat_type"] = chatType - } - if sender != nil && sender.TenantKey != nil { - metadata["tenant_key"] = *sender.TenantKey - } + metadata := buildInboundMetadata(message, sender) var peer bus.Peer if chatType == "p2p" { @@ -477,12 +468,25 @@ func (c *FeishuChannel) handleMessageReceive(ctx context.Context, event *larkim. content = cleaned } + if replyTargetID(message) != "" || stringValue(message.ThreadId) != "" { + content, mediaRefs = c.prependReplyContext(ctx, message, chatID, content, mediaRefs) + } + if content == "" { + content = "[empty message]" + } + logger.InfoCF("feishu", "Feishu message received", map[string]any{ "sender_id": senderID, "chat_id": chatID, "message_id": messageID, "preview": utils.Truncate(content, 80), }) + logger.InfoCF("feishu", "Feishu reply linkage", map[string]any{ + "message_id": messageID, + "parent_id": stringValue(message.ParentId), + "root_id": stringValue(message.RootId), + "thread_id": stringValue(message.ThreadId), + }) c.HandleMessage(ctx, peer, messageID, senderID, chatID, content, mediaRefs, metadata, senderInfo) return nil diff --git a/pkg/channels/feishu/feishu_reply.go b/pkg/channels/feishu/feishu_reply.go new file mode 100644 index 000000000..22dfe3e87 --- /dev/null +++ b/pkg/channels/feishu/feishu_reply.go @@ -0,0 +1,298 @@ +//go:build amd64 || arm64 || riscv64 || mips64 || ppc64 + +package feishu + +import ( + "context" + "fmt" + "strings" + "time" + + larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1" + + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/utils" +) + +const messageCacheTTL = 30 * time.Second + +const ( + maxReplyContextLen = 600 +) + +func (c *FeishuChannel) prependReplyContext( + ctx context.Context, + message *larkim.EventMessage, + chatID string, + content string, + mediaRefs []string, +) (string, []string) { + if message == nil { + return content, mediaRefs + } + + lookupCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + targetMessageID := c.resolveReplyTargetMessageID(lookupCtx, message) + if targetMessageID == "" { + logger.DebugCF("feishu", "No reply target resolved; skip reply context", map[string]any{ + "message_id": stringValue(message.MessageId), + "parent_id": stringValue(message.ParentId), + "root_id": stringValue(message.RootId), + "thread_id": stringValue(message.ThreadId), + }) + return content, mediaRefs + } + + repliedMessage, err := c.fetchMessageByID(lookupCtx, targetMessageID) + if err != nil { + logger.DebugCF("feishu", "Failed to fetch replied message context", map[string]any{ + "target_message_id": targetMessageID, + "error": err.Error(), + }) + return content, mediaRefs + } + + messageType := stringValue(repliedMessage.MsgType) + rawContent := "" + if repliedMessage.Body != nil { + rawContent = stringValue(repliedMessage.Body.Content) + } + + var repliedMediaRefs []string + if store := c.GetMediaStore(); store != nil { + repliedMediaRefs = c.downloadInboundMedia(lookupCtx, chatID, targetMessageID, messageType, rawContent, store) + if messageType == larkim.MsgTypeInteractive { + _, externalURLs := extractCardImageKeys(rawContent) + if len(externalURLs) > 0 { + repliedMediaRefs = append(repliedMediaRefs, externalURLs...) + } + } + } + + repliedContent := normalizeRepliedContent(messageType, rawContent, repliedMediaRefs) + if len(repliedMediaRefs) > 0 { + mediaRefs = append(repliedMediaRefs, mediaRefs...) + } + + return formatReplyContext(targetMessageID, repliedContent, content), mediaRefs +} + +func (c *FeishuChannel) resolveReplyTargetMessageID(ctx context.Context, message *larkim.EventMessage) string { + if targetID := replyTargetID(message); targetID != "" { + logger.DebugCF("feishu", "Resolved reply target from event payload", map[string]any{ + "message_id": stringValue(message.MessageId), + "parent_id": stringValue(message.ParentId), + "root_id": stringValue(message.RootId), + "target_id": targetID, + }) + return targetID + } + + currentMessageID := stringValue(message.MessageId) + if currentMessageID == "" { + return "" + } + + if stringValue(message.ThreadId) == "" { + logger.DebugCF("feishu", "No reply target found; message is not in a thread", map[string]any{ + "message_id": stringValue(message.MessageId), + }) + return "" + } + + msg, err := c.fetchMessageByID(ctx, currentMessageID) + if err != nil { + logger.DebugCF("feishu", "Failed to query current message detail for reply info", map[string]any{ + "message_id": currentMessageID, + "error": err.Error(), + }) + return "" + } + + targetID := replyTargetIDFromMessage(msg) + if targetID != "" { + logger.DebugCF("feishu", "Resolved reply target from message detail", map[string]any{ + "message_id": currentMessageID, + "parent_id": stringValue(msg.ParentId), + "root_id": stringValue(msg.RootId), + "target_id": targetID, + }) + } + return targetID +} + +func (c *FeishuChannel) fetchMessageByID(ctx context.Context, messageID string) (*larkim.Message, error) { + if cached, ok := c.messageCache.Load(messageID); ok { + cm := cached.(*cachedMessage) + if time.Now().Before(cm.expiry) { + return cm.msg, nil + } + c.messageCache.Delete(messageID) + } + + req := larkim.NewGetMessageReqBuilder(). + MessageId(messageID). + Build() + + resp, err := c.client.Im.V1.Message.Get(ctx, req) + if err != nil { + return nil, fmt.Errorf("feishu get message: %w", err) + } + if !resp.Success() { + c.invalidateTokenOnAuthError(resp.Code) + return nil, fmt.Errorf("feishu get message api error (code=%d msg=%s)", resp.Code, resp.Msg) + } + if resp.Data == nil || len(resp.Data.Items) == 0 || resp.Data.Items[0] == nil { + return nil, fmt.Errorf("feishu get message: empty response") + } + // Items[0] contains the target message - the Feishu API returns a list + // but we request a single message by ID, so the list always has at most one item. + msg := resp.Data.Items[0] + c.messageCache.Store(messageID, &cachedMessage{msg: msg, expiry: time.Now().Add(messageCacheTTL)}) + return msg, nil +} + +func replyTargetID(message *larkim.EventMessage) string { + if message == nil { + return "" + } + if parentID := stringValue(message.ParentId); parentID != "" { + return parentID + } + return stringValue(message.RootId) +} + +func replyTargetIDFromMessage(message *larkim.Message) string { + if message == nil { + return "" + } + if parentID := stringValue(message.ParentId); parentID != "" { + return parentID + } + return stringValue(message.RootId) +} + +func buildInboundMetadata(message *larkim.EventMessage, sender *larkim.EventSender) map[string]string { + metadata := map[string]string{} + if message == nil { + return metadata + } + + messageID := stringValue(message.MessageId) + if messageID != "" { + metadata["message_id"] = messageID + } + + messageType := stringValue(message.MessageType) + if messageType != "" { + metadata["message_type"] = messageType + } + + chatType := stringValue(message.ChatType) + if chatType != "" { + metadata["chat_type"] = chatType + } + + parentID := stringValue(message.ParentId) + if parentID != "" { + metadata["parent_id"] = parentID + } + + rootID := stringValue(message.RootId) + if rootID != "" { + metadata["root_id"] = rootID + } + + if replyTo := replyTargetID(message); replyTo != "" { + metadata["reply_to_message_id"] = replyTo + } + + threadID := stringValue(message.ThreadId) + if threadID != "" { + metadata["thread_id"] = threadID + } + + if sender != nil && sender.TenantKey != nil && *sender.TenantKey != "" { + metadata["tenant_key"] = *sender.TenantKey + } + + return metadata +} + +func normalizeRepliedContent(messageType, rawContent string, mediaRefs []string) string { + content := extractContent(messageType, rawContent) + + if containsFeishuUpgradePlaceholder(rawContent) || containsFeishuUpgradePlaceholder(content) { + content = "" + } + + content = appendMediaTags(content, messageType, mediaRefs) + if strings.TrimSpace(content) != "" { + return content + } + + switch messageType { + case larkim.MsgTypeImage: + return "[replied image]" + case larkim.MsgTypeFile: + return "[replied file]" + case larkim.MsgTypeAudio: + return "[replied audio]" + case larkim.MsgTypeMedia: + return "[replied video]" + case larkim.MsgTypeInteractive: + return "[replied interactive card]" + default: + return "[replied message content unavailable]" + } +} + +func containsFeishuUpgradePlaceholder(s string) bool { + upgradePrompt := "\u8bf7\u5347\u7ea7\u81f3\u6700\u65b0\u7248\u672c\u5ba2\u6237\u7aef" + upgradePromptEscaped := "\\u8bf7\\u5347\\u7ea7\\u81f3\\u6700\\u65b0\\u7248\\u672c\\u5ba2\\u6237\\u7aef" + return strings.Contains(s, upgradePrompt) || strings.Contains(s, upgradePromptEscaped) +} + +func formatReplyContext(parentID, repliedContent, content string) string { + parentID = strings.TrimSpace(parentID) + repliedContent = strings.TrimSpace(repliedContent) + content = strings.TrimSpace(content) + + if parentID == "" || repliedContent == "" { + return content + } + + repliedContent = utils.Truncate(repliedContent, maxReplyContextLen) + repliedContent = sanitizeReplyContextContent(repliedContent) + content = sanitizeReplyContextContent(content) + header := fmt.Sprintf("[replied_message id=%q]", parentID) + footer := "[/replied_message]" + if content == "" { + return header + "\n" + repliedContent + "\n" + footer + } + if hasLeadingCommandPrefix(content) { + return content + "\n\n" + header + "\n" + repliedContent + "\n" + footer + } + return header + "\n" + repliedContent + "\n" + footer + "\n\n[current_message]\n" + content + "\n[/current_message]" +} + +func hasLeadingCommandPrefix(s string) bool { + tokens := strings.Fields(strings.TrimSpace(s)) + if len(tokens) == 0 { + return false + } + first := tokens[0] + return strings.HasPrefix(first, "/") || strings.HasPrefix(first, "!") +} + +func sanitizeReplyContextContent(s string) string { + tagEscaper := strings.NewReplacer( + "[replied_message", `\[replied_message`, + "[/replied_message]", `\[/replied_message]`, + "[current_message]", `\[current_message]`, + "[/current_message]", `\[/current_message]`, + ) + return tagEscaper.Replace(s) +} diff --git a/pkg/channels/feishu/feishu_reply_test.go b/pkg/channels/feishu/feishu_reply_test.go new file mode 100644 index 000000000..0efe7bc01 --- /dev/null +++ b/pkg/channels/feishu/feishu_reply_test.go @@ -0,0 +1,229 @@ +//go:build amd64 || arm64 || riscv64 || mips64 || ppc64 + +package feishu + +import ( + "strings" + "testing" + + larkim "github.com/larksuite/oapi-sdk-go/v3/service/im/v1" +) + +func TestBuildInboundMetadata(t *testing.T) { + strPtr := func(s string) *string { return &s } + + t.Run("includes basic and reply fields", func(t *testing.T) { + message := &larkim.EventMessage{ + MessageId: strPtr("om_msg_1"), + MessageType: strPtr("text"), + ChatType: strPtr("group"), + ParentId: strPtr("om_parent_1"), + RootId: strPtr("om_root_1"), + ThreadId: strPtr("omt_thread_1"), + } + sender := &larkim.EventSender{TenantKey: strPtr("tenant_x")} + + got := buildInboundMetadata(message, sender) + + if got["message_id"] != "om_msg_1" { + t.Fatalf("message_id = %q, want %q", got["message_id"], "om_msg_1") + } + if got["message_type"] != "text" { + t.Fatalf("message_type = %q, want %q", got["message_type"], "text") + } + if got["chat_type"] != "group" { + t.Fatalf("chat_type = %q, want %q", got["chat_type"], "group") + } + if got["parent_id"] != "om_parent_1" { + t.Fatalf("parent_id = %q, want %q", got["parent_id"], "om_parent_1") + } + if got["reply_to_message_id"] != "om_parent_1" { + t.Fatalf("reply_to_message_id = %q, want %q", got["reply_to_message_id"], "om_parent_1") + } + if got["root_id"] != "om_root_1" { + t.Fatalf("root_id = %q, want %q", got["root_id"], "om_root_1") + } + if got["thread_id"] != "omt_thread_1" { + t.Fatalf("thread_id = %q, want %q", got["thread_id"], "omt_thread_1") + } + if got["tenant_key"] != "tenant_x" { + t.Fatalf("tenant_key = %q, want %q", got["tenant_key"], "tenant_x") + } + }) + + t.Run("falls back reply_to_message_id to root_id", func(t *testing.T) { + message := &larkim.EventMessage{ + MessageId: strPtr("om_msg_3"), + RootId: strPtr("om_root_3"), + } + + got := buildInboundMetadata(message, nil) + + if got["root_id"] != "om_root_3" { + t.Fatalf("root_id = %q, want %q", got["root_id"], "om_root_3") + } + if got["reply_to_message_id"] != "om_root_3" { + t.Fatalf("reply_to_message_id = %q, want %q", got["reply_to_message_id"], "om_root_3") + } + }) + + t.Run("omits empty values", func(t *testing.T) { + message := &larkim.EventMessage{ + MessageId: strPtr("om_msg_2"), + } + + got := buildInboundMetadata(message, nil) + + if got["message_id"] != "om_msg_2" { + t.Fatalf("message_id = %q, want %q", got["message_id"], "om_msg_2") + } + if _, ok := got["parent_id"]; ok { + t.Fatalf("parent_id should be absent, got %q", got["parent_id"]) + } + if _, ok := got["reply_to_message_id"]; ok { + t.Fatalf("reply_to_message_id should be absent, got %q", got["reply_to_message_id"]) + } + if _, ok := got["tenant_key"]; ok { + t.Fatalf("tenant_key should be absent, got %q", got["tenant_key"]) + } + }) + + t.Run("nil message returns empty map", func(t *testing.T) { + got := buildInboundMetadata(nil, nil) + if len(got) != 0 { + t.Fatalf("len(metadata) = %d, want 0", len(got)) + } + }) +} + +func TestFormatReplyContext(t *testing.T) { + t.Run("formats reply context with content", func(t *testing.T) { + got := formatReplyContext("om_parent_1", "original message", "new reply") + want := "[replied_message id=\"om_parent_1\"]\noriginal message\n[/replied_message]\n\n[current_message]\nnew reply\n[/current_message]" + if got != want { + t.Fatalf("formatReplyContext() = %q, want %q", got, want) + } + }) + + t.Run("returns reply context when current content is empty", func(t *testing.T) { + got := formatReplyContext("om_parent_1", "original message", "") + want := "[replied_message id=\"om_parent_1\"]\noriginal message\n[/replied_message]" + if got != want { + t.Fatalf("formatReplyContext() = %q, want %q", got, want) + } + }) + + t.Run("returns original content when parent or replied content missing", func(t *testing.T) { + if got := formatReplyContext("", "original", "new reply"); got != "new reply" { + t.Fatalf("missing parent: got %q, want %q", got, "new reply") + } + if got := formatReplyContext("om_parent_1", "", "new reply"); got != "new reply" { + t.Fatalf("missing replied content: got %q, want %q", got, "new reply") + } + }) + + t.Run("escapes reserved wrapper tags in payload", func(t *testing.T) { + replied := "payload [replied_message id=\"x\"] x [/replied_message]" + current := "hello [current_message]injected[/current_message]" + got := formatReplyContext("om_parent_1", replied, current) + + if !strings.HasPrefix(got, "[replied_message id=\"om_parent_1\"]") { + t.Fatalf("outer replied_message wrapper missing: %q", got) + } + if strings.Contains(got, "\n[replied_message id=\"x\"]") { + t.Fatalf("nested replied_message tag should be escaped: %q", got) + } + if strings.Contains(got, "\n[current_message]injected") { + t.Fatalf("nested current_message tag should be escaped: %q", got) + } + if !strings.Contains(got, `\[replied_message id="x"]`) { + t.Fatalf("escaped replied tag missing: %q", got) + } + }) + + t.Run("preserves leading slash command prefix", func(t *testing.T) { + got := formatReplyContext("om_parent_1", "original message", "/help") + want := "/help\n\n[replied_message id=\"om_parent_1\"]\noriginal message\n[/replied_message]" + if got != want { + t.Fatalf("formatReplyContext() = %q, want %q", got, want) + } + }) + + t.Run("preserves leading bang command prefix", func(t *testing.T) { + got := formatReplyContext("om_parent_1", "original message", "!status now") + want := "!status now\n\n[replied_message id=\"om_parent_1\"]\noriginal message\n[/replied_message]" + if got != want { + t.Fatalf("formatReplyContext() = %q, want %q", got, want) + } + }) +} + +func TestReplyTargetID(t *testing.T) { + strPtr := func(s string) *string { return &s } + + t.Run("prefer parent_id", func(t *testing.T) { + msg := &larkim.EventMessage{ParentId: strPtr("om_parent"), RootId: strPtr("om_root")} + if got := replyTargetID(msg); got != "om_parent" { + t.Fatalf("replyTargetID() = %q, want %q", got, "om_parent") + } + }) + + t.Run("fallback to root_id", func(t *testing.T) { + msg := &larkim.EventMessage{RootId: strPtr("om_root")} + if got := replyTargetID(msg); got != "om_root" { + t.Fatalf("replyTargetID() = %q, want %q", got, "om_root") + } + }) + + t.Run("empty when no fields", func(t *testing.T) { + if got := replyTargetID(&larkim.EventMessage{}); got != "" { + t.Fatalf("replyTargetID() = %q, want empty", got) + } + }) +} + +func TestNormalizeRepliedContent(t *testing.T) { + t.Run("filters feishu upgrade placeholder for interactive", func(t *testing.T) { + raw := `{"text":"\u8bf7\u5347\u7ea7\u81f3\u6700\u65b0\u7248\u672c\u5ba2\u6237\u7aef\uff0c\u4ee5\u67e5\u770b\u5185\u5bb9"}` + got := normalizeRepliedContent("interactive", raw, nil) + if got != "[replied interactive card]" { + t.Fatalf("normalizeRepliedContent() = %q, want %q", got, "[replied interactive card]") + } + }) + + t.Run("keeps filename and file tag for replied file", func(t *testing.T) { + got := normalizeRepliedContent("file", `{"file_key":"file_xxx","file_name":"doc.pdf"}`, []string{"media://r1"}) + if got != "doc.pdf [file]" { + t.Fatalf("normalizeRepliedContent() = %q, want %q", got, "doc.pdf [file]") + } + }) + + t.Run("falls back when file content missing", func(t *testing.T) { + got := normalizeRepliedContent("file", `{"file_key":"file_xxx"}`, nil) + if got != "[replied file]" { + t.Fatalf("normalizeRepliedContent() = %q, want %q", got, "[replied file]") + } + }) +} + +func TestHasLeadingCommandPrefix(t *testing.T) { + tests := []struct { + name string + input string + want bool + }{ + {name: "slash command", input: "/help", want: true}, + {name: "bang command", input: "!status", want: true}, + {name: "leading spaces slash", input: " /ping arg", want: true}, + {name: "normal text", input: "hello /help", want: false}, + {name: "empty", input: "", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := hasLeadingCommandPrefix(tt.input); got != tt.want { + t.Fatalf("hasLeadingCommandPrefix(%q) = %v, want %v", tt.input, got, tt.want) + } + }) + } +} diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index 6d9f5eda8..c4326fda0 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -430,6 +430,19 @@ func (m *Manager) initChannels(channels *config.ChannelsConfig) error { m.initChannel("vk", "VK") } + if channels.TeamsWebhook.Enabled && len(channels.TeamsWebhook.Webhooks) > 0 { + hasValidTarget := false + for _, target := range channels.TeamsWebhook.Webhooks { + if target.WebhookURL.String() != "" { + hasValidTarget = true + break + } + } + if hasValidTarget { + m.initChannel("teams_webhook", "Teams Webhook") + } + } + logger.InfoCF("channels", "Channel initialization completed", map[string]any{ "enabled_channels": len(m.channels), }) diff --git a/pkg/channels/manager_channel.go b/pkg/channels/manager_channel.go index b1c8c25e0..b54facda4 100644 --- a/pkg/channels/manager_channel.go +++ b/pkg/channels/manager_channel.go @@ -62,6 +62,13 @@ func hiddenValues(key string, value map[string]any, ch config.ChannelsConfig) { value["app_secret"] = ch.Feishu.AppSecret.String() value["encrypt_key"] = ch.Feishu.EncryptKey.String() value["verification_token"] = ch.Feishu.VerificationToken.String() + case "teams_webhook": + // Expose webhook URLs for hash computation (they contain secrets) + webhooks := make(map[string]string) + for name, target := range ch.TeamsWebhook.Webhooks { + webhooks[name] = target.WebhookURL.String() + } + value["webhooks"] = webhooks } } @@ -166,4 +173,13 @@ func updateKeys(newcfg, old *config.ChannelsConfig) { newcfg.Feishu.EncryptKey = old.Feishu.EncryptKey newcfg.Feishu.VerificationToken = old.Feishu.VerificationToken } + if newcfg.TeamsWebhook.Enabled { + // Copy SecureString webhook URLs from old config + for name, oldTarget := range old.TeamsWebhook.Webhooks { + if newTarget, ok := newcfg.TeamsWebhook.Webhooks[name]; ok { + newTarget.WebhookURL = oldTarget.WebhookURL + newcfg.TeamsWebhook.Webhooks[name] = newTarget + } + } + } } diff --git a/pkg/channels/teams_webhook/init.go b/pkg/channels/teams_webhook/init.go new file mode 100644 index 000000000..fca960039 --- /dev/null +++ b/pkg/channels/teams_webhook/init.go @@ -0,0 +1,13 @@ +package teamswebhook + +import ( + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels" + "github.com/sipeed/picoclaw/pkg/config" +) + +func init() { + channels.RegisterFactory("teams_webhook", func(cfg *config.Config, b *bus.MessageBus) (channels.Channel, error) { + return NewTeamsWebhookChannel(cfg.Channels.TeamsWebhook, b) + }) +} diff --git a/pkg/channels/teams_webhook/teams_webhook.go b/pkg/channels/teams_webhook/teams_webhook.go new file mode 100644 index 000000000..fa7762a3e --- /dev/null +++ b/pkg/channels/teams_webhook/teams_webhook.go @@ -0,0 +1,422 @@ +package teamswebhook + +import ( + "context" + "fmt" + "net/url" + "regexp" + "sort" + "strconv" + "strings" + + goteamsnotify "github.com/atc0005/go-teams-notify/v2" + "github.com/atc0005/go-teams-notify/v2/adaptivecard" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +// statusCodeRe extracts HTTP status codes from error messages like "401 Unauthorized". +var statusCodeRe = regexp.MustCompile(`\b([45]\d{2})\b`) + +// markdownTableRe matches a markdown table block (header + separator + rows). +// It captures the entire table including all rows. +var markdownTableRe = regexp.MustCompile(`(?m)^(\|[^\n]+\|)\n(\|[-:\|\s]+\|)\n((?:\|[^\n]+\|\n?)+)`) + +// teamsMessageSender abstracts the Teams client for testability. +type teamsMessageSender interface { + SendWithContext(ctx context.Context, webhookURL string, message goteamsnotify.TeamsMessage) error +} + +// classifyTeamsError extracts HTTP status code from error message and classifies it. +// The go-teams-notify library returns errors like "error on notification: 401 Unauthorized, ...". +// This allows proper retry behavior: 4xx errors are permanent, 5xx are temporary. +func classifyTeamsError(err error) error { + if err == nil { + return nil + } + errMsg := err.Error() + if matches := statusCodeRe.FindStringSubmatch(errMsg); len(matches) > 1 { + if statusCode, parseErr := strconv.Atoi(matches[1]); parseErr == nil { + return channels.ClassifySendError(statusCode, err) + } + } + // Fallback: treat as temporary network error (retryable) + return channels.ClassifyNetError(err) +} + +// TeamsWebhookChannel is an output-only channel that sends messages +// to Microsoft Teams via Power Automate workflow webhooks. +// Multiple webhook targets can be configured and selected via ChatID. +type TeamsWebhookChannel struct { + *channels.BaseChannel + config config.TeamsWebhookConfig + client teamsMessageSender +} + +// NewTeamsWebhookChannel creates a new Teams webhook channel. +func NewTeamsWebhookChannel( + cfg config.TeamsWebhookConfig, + bus *bus.MessageBus, +) (*TeamsWebhookChannel, error) { + if len(cfg.Webhooks) == 0 { + return nil, fmt.Errorf("teams_webhook: at least one webhook target is required") + } + + // Require "default" webhook target + if _, hasDefault := cfg.Webhooks["default"]; !hasDefault { + return nil, fmt.Errorf("teams_webhook: a 'default' webhook target is required") + } + + // Validate all webhook targets have valid HTTPS URLs + for name, target := range cfg.Webhooks { + webhookURL := target.WebhookURL.String() + if webhookURL == "" { + return nil, fmt.Errorf("teams_webhook: webhook %q has empty webhook_url", name) + } + parsed, err := url.Parse(webhookURL) + if err != nil { + return nil, fmt.Errorf("teams_webhook: webhook %q has invalid URL: %w", name, err) + } + if !strings.EqualFold(parsed.Scheme, "https") { + return nil, fmt.Errorf("teams_webhook: webhook %q must use HTTPS (got %q)", name, parsed.Scheme) + } + } + + base := channels.NewBaseChannel( + "teams_webhook", + cfg, + bus, + []string{ + "*", + }, // Output-only channel; "*" suppresses misleading "allows EVERYONE" audit warning + channels.WithMaxMessageLength(24000), // Power Automate webhook payload limit is 28KB + ) + + client := goteamsnotify.NewTeamsClient() + + return &TeamsWebhookChannel{ + BaseChannel: base, + config: cfg, + client: client, + }, nil +} + +// Start initializes the channel. For output-only channels, this is a no-op. +func (c *TeamsWebhookChannel) Start(ctx context.Context) error { + targets := make([]string, 0, len(c.config.Webhooks)) + for name := range c.config.Webhooks { + targets = append(targets, name) + } + sort.Strings(targets) + logger.InfoCF("teams_webhook", "Starting Teams webhook channel (output-only)", map[string]any{ + "targets": targets, + }) + c.SetRunning(true) + return nil +} + +// Stop shuts down the channel. +func (c *TeamsWebhookChannel) Stop(ctx context.Context) error { + logger.InfoC("teams_webhook", "Stopping Teams webhook channel") + c.SetRunning(false) + return nil +} + +// Send delivers a message to the specified Teams webhook target. +// The target is selected by msg.ChatID which must match a key in the webhooks map. +func (c *TeamsWebhookChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]string, error) { + if !c.IsRunning() { + return nil, channels.ErrNotRunning + } + + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + // Look up webhook target by ChatID, fall back to "default" if empty or unknown + targetName := msg.ChatID + if targetName == "" { + targetName = "default" + } + + target, ok := c.config.Webhooks[targetName] + if !ok { + // Log warning and fall back to default target + logger.WarnCF("teams_webhook", "Unknown target, falling back to default", map[string]any{ + "requested": msg.ChatID, + "using": "default", + }) + target = c.config.Webhooks["default"] + } + + // Build an Adaptive Card for rich formatting + card, err := c.buildAdaptiveCard(msg, target) + if err != nil { + return nil, fmt.Errorf("teams_webhook: failed to build card: %w", err) + } + + // Create the message with the card + teamsMsg, err := adaptivecard.NewMessageFromCard(card) + if err != nil { + return nil, fmt.Errorf("teams_webhook: failed to create message: %w", err) + } + + // Send to Teams + if err := c.client.SendWithContext(ctx, target.WebhookURL.String(), teamsMsg); err != nil { + // Log without raw error to avoid leaking webhook URL (embedded in net/http errors) + logger.ErrorCF("teams_webhook", "Failed to send message to Teams webhook", map[string]any{ + "target": msg.ChatID, + }) + // Classify error based on status code extracted from error message. + // The go-teams-notify library includes status in errors like "401 Unauthorized". + // Use ClassifySendError for proper retry behavior (4xx = permanent, 5xx = temporary). + classifiedErr := classifyTeamsError(err) + return nil, fmt.Errorf("teams_webhook: send failed: %w", classifiedErr) + } + + logger.DebugCF("teams_webhook", "Message sent successfully", map[string]any{ + "target": msg.ChatID, + }) + + return nil, nil +} + +// buildAdaptiveCard creates a formatted Adaptive Card from the outbound message. +// It detects markdown tables and converts them to native Adaptive Card Table elements, +// since TextBlocks only support a limited markdown subset (no tables). +func (c *TeamsWebhookChannel) buildAdaptiveCard( + msg bus.OutboundMessage, + target config.TeamsWebhookTarget, +) (adaptivecard.Card, error) { + card := adaptivecard.NewCard() + card.Type = adaptivecard.TypeAdaptiveCard + + // Set full width for Teams rendering + card.MSTeams.Width = "Full" + + // Add title if configured on the target + title := target.Title + if title == "" { + title = "PicoClaw Notification" + } + + titleBlock := adaptivecard.NewTextBlock(title, true) + titleBlock.Size = adaptivecard.SizeLarge + titleBlock.Weight = adaptivecard.WeightBolder + titleBlock.Style = adaptivecard.TextBlockStyleHeading + + if err := card.AddElement(false, titleBlock); err != nil { + return card, err + } + + content := msg.Content + if content == "" { + content = "(empty message)" + } + + // Split content into text segments and tables + // TextBlocks support: bold, italic, bullet/numbered lists, links + // TextBlocks do NOT support: headers, tables, images + segments := splitContentWithTables(content) + + for _, seg := range segments { + if seg.isTable { + // Convert markdown table to Adaptive Card Table element + tableElement, err := parseMarkdownTable(seg.content) + if err != nil { + // Fallback: render as preformatted text if parsing fails + logger.WarnCF("teams_webhook", "Failed to parse markdown table, using fallback", map[string]any{ + "error": err.Error(), + }) + block := adaptivecard.NewTextBlock("```\n"+seg.content+"\n```", true) + block.Wrap = true + if err := card.AddElement(false, block); err != nil { + return card, err + } + continue + } + if err := card.AddElement(false, tableElement); err != nil { + return card, err + } + } else { + // Regular text content + text := strings.TrimSpace(seg.content) + if text == "" { + continue + } + block := adaptivecard.NewTextBlock(text, true) + block.Wrap = true + if err := card.AddElement(false, block); err != nil { + return card, err + } + } + } + + return card, nil +} + +// contentSegment represents either a text block or a table in the message content. +type contentSegment struct { + content string + isTable bool +} + +// splitContentWithTables splits content into alternating text and table segments. +func splitContentWithTables(content string) []contentSegment { + var segments []contentSegment + + matches := markdownTableRe.FindAllStringSubmatchIndex(content, -1) + if len(matches) == 0 { + // No tables found, return entire content as text + return []contentSegment{{content: content, isTable: false}} + } + + lastEnd := 0 + for _, match := range matches { + // Text before this table + if match[0] > lastEnd { + segments = append(segments, contentSegment{ + content: content[lastEnd:match[0]], + isTable: false, + }) + } + // The table itself + segments = append(segments, contentSegment{ + content: content[match[0]:match[1]], + isTable: true, + }) + lastEnd = match[1] + } + + // Text after the last table + if lastEnd < len(content) { + segments = append(segments, contentSegment{ + content: content[lastEnd:], + isTable: false, + }) + } + + return segments +} + +// parseMarkdownTable converts a markdown table string to an Adaptive Card Table element. +func parseMarkdownTable(tableStr string) (adaptivecard.Element, error) { + lines := strings.Split(strings.TrimSpace(tableStr), "\n") + if len(lines) < 2 { + return adaptivecard.Element{}, fmt.Errorf("table must have at least header and separator rows") + } + + // Track header content length per column for width calculation + var headerLengths []int + + // Parse all rows (header + data rows, skip separator) + var allRows [][]adaptivecard.TableCell + for i, line := range lines { + // Skip separator row (contains only |, -, :, and spaces) + if i == 1 && isSeparatorRow(line) { + continue + } + + cells := parseTableRow(line) + if len(cells) == 0 { + continue + } + + var tableCells []adaptivecard.TableCell + for _, cellText := range cells { + trimmedText := strings.TrimSpace(cellText) + + // Use header row (first row) to determine column widths + if i == 0 { + headerLengths = append(headerLengths, len(trimmedText)) + } + + textBlock := adaptivecard.Element{ + Type: adaptivecard.TypeElementTextBlock, + Text: trimmedText, + Wrap: true, + } + cell := adaptivecard.TableCell{ + Type: adaptivecard.TypeTableCell, + Items: []*adaptivecard.Element{&textBlock}, + } + tableCells = append(tableCells, cell) + } + allRows = append(allRows, tableCells) + } + + if len(allRows) == 0 { + return adaptivecard.Element{}, fmt.Errorf("no valid rows found in table") + } + + // Create table with first row as headers + firstRowAsHeaders := true + showGridLines := true + + table, err := adaptivecard.NewTableFromTableCells(allRows, 0, firstRowAsHeaders, showGridLines) + if err != nil { + return adaptivecard.Element{}, fmt.Errorf("failed to create table: %w", err) + } + + // Set column widths based on header content length + table.Columns = calculateColumnWidths(headerLengths) + + return table, nil +} + +// calculateColumnWidths creates TableColumnDefinition entries with widths +// proportional to the max content length of each column. +func calculateColumnWidths(maxLengths []int) []adaptivecard.Column { + if len(maxLengths) == 0 { + return nil + } + + // Use content length as relative weight, with a minimum of 1 + columns := make([]adaptivecard.Column, len(maxLengths)) + for i, length := range maxLengths { + weight := length + if weight < 1 { + weight = 1 + } + columns[i] = adaptivecard.Column{ + Type: "TableColumnDefinition", + Width: weight, + } + } + + return columns +} + +// isSeparatorRow checks if a line is a markdown table separator (e.g., |---|---|). +func isSeparatorRow(line string) bool { + // Remove pipes and spaces, check if only dashes and colons remain + cleaned := strings.ReplaceAll(line, "|", "") + cleaned = strings.ReplaceAll(cleaned, " ", "") + cleaned = strings.ReplaceAll(cleaned, "-", "") + cleaned = strings.ReplaceAll(cleaned, ":", "") + return cleaned == "" +} + +// parseTableRow extracts cell values from a markdown table row. +func parseTableRow(line string) []string { + // Trim leading/trailing pipes and split by | + line = strings.TrimSpace(line) + line = strings.TrimPrefix(line, "|") + line = strings.TrimSuffix(line, "|") + + if line == "" { + return nil + } + + parts := strings.Split(line, "|") + var cells []string + for _, p := range parts { + cells = append(cells, strings.TrimSpace(p)) + } + return cells +} diff --git a/pkg/channels/teams_webhook/teams_webhook_test.go b/pkg/channels/teams_webhook/teams_webhook_test.go new file mode 100644 index 000000000..451ba9d18 --- /dev/null +++ b/pkg/channels/teams_webhook/teams_webhook_test.go @@ -0,0 +1,583 @@ +package teamswebhook + +import ( + "context" + "errors" + "testing" + + goteamsnotify "github.com/atc0005/go-teams-notify/v2" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" +) + +// mockTeamsClient implements teamsMessageSender for testing. +type mockTeamsClient struct { + sendFunc func(ctx context.Context, webhookURL string, message goteamsnotify.TeamsMessage) error +} + +func (m *mockTeamsClient) SendWithContext( + ctx context.Context, + webhookURL string, + message goteamsnotify.TeamsMessage, +) error { + if m.sendFunc != nil { + return m.sendFunc(ctx, webhookURL, message) + } + return nil +} + +func TestNewTeamsWebhookChannel(t *testing.T) { + msgBus := bus.NewMessageBus() + + // Test missing webhooks + _, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: nil, + }, msgBus) + if err == nil { + t.Error("expected error for missing webhooks") + } + + // Test missing "default" webhook + _, err = NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "alerts": { + WebhookURL: *config.NewSecureString("https://example.com/webhook"), + Title: "Alerts", + }, + }, + }, msgBus) + if err == nil { + t.Error("expected error for missing 'default' webhook") + } + + // Test empty webhook URL + _, err = NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": {Title: "Default"}, + }, + }, msgBus) + if err == nil { + t.Error("expected error for empty webhook_url") + } + + // Test HTTP URL (should fail, must be HTTPS) + _, err = NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("http://example.com/webhook"), + Title: "Default", + }, + }, + }, msgBus) + if err == nil { + t.Error("expected error for HTTP webhook URL (must be HTTPS)") + } + + // Test valid config with HTTPS (must include "default") + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-default"), + Title: "Default", + }, + "alerts": { + WebhookURL: *config.NewSecureString("https://example.com/webhook1"), + Title: "Alerts", + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if ch.Name() != "teams_webhook" { + t.Errorf("expected name 'teams_webhook', got %q", ch.Name()) + } +} + +func TestTeamsWebhookChannel_StartStop(t *testing.T) { + msgBus := bus.NewMessageBus() + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook"), + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx := context.Background() + + if ch.IsRunning() { + t.Error("channel should not be running before Start") + } + + if err := ch.Start(ctx); err != nil { + t.Fatalf("Start failed: %v", err) + } + + if !ch.IsRunning() { + t.Error("channel should be running after Start") + } + + if err := ch.Stop(ctx); err != nil { + t.Fatalf("Stop failed: %v", err) + } + + if ch.IsRunning() { + t.Error("channel should not be running after Stop") + } +} + +func TestTeamsWebhookChannel_BuildAdaptiveCard(t *testing.T) { + msgBus := bus.NewMessageBus() + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-default"), + Title: "Default", + }, + "alerts": { + WebhookURL: *config.NewSecureString("https://example.com/webhook"), + Title: "Custom Title", + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + target := ch.config.Webhooks["alerts"] + msg := bus.OutboundMessage{ + Content: "Test message content", + ChatID: "alerts", + } + + card, err := ch.buildAdaptiveCard(msg, target) + if err != nil { + t.Fatalf("buildAdaptiveCard failed: %v", err) + } + + if card.Type != "AdaptiveCard" { + t.Errorf("expected card type 'AdaptiveCard', got %q", card.Type) + } +} + +func TestTeamsWebhookChannel_SendNotRunning(t *testing.T) { + msgBus := bus.NewMessageBus() + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook"), + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx := context.Background() + msg := bus.OutboundMessage{Content: "test", ChatID: "default"} + + _, err = ch.Send(ctx, msg) + if err == nil { + t.Error("expected error when sending while not running") + } +} + +func TestTeamsWebhookChannel_SendDefaultTargetFallback(t *testing.T) { + tests := []struct { + name string + chatID string + }{ + {"unknown target falls back to default", "unknown"}, + {"empty ChatID uses default", ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + msgBus := bus.NewMessageBus() + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-default"), + }, + "alerts": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-alerts"), + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var sentURL string + ch.client = &mockTeamsClient{ + sendFunc: func(ctx context.Context, webhookURL string, message goteamsnotify.TeamsMessage) error { + sentURL = webhookURL + return nil + }, + } + + ctx := context.Background() + _ = ch.Start(ctx) + defer ch.Stop(ctx) + + msg := bus.OutboundMessage{Content: "test", ChatID: tt.chatID} + _, err = ch.Send(ctx, msg) + if err != nil { + t.Fatalf("expected success, got error: %v", err) + } + + if sentURL != "https://example.com/webhook-default" { + t.Errorf("expected default webhook URL, got %q", sentURL) + } + }) + } +} + +func TestTeamsWebhookChannel_SendSuccess(t *testing.T) { + msgBus := bus.NewMessageBus() + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-default"), + Title: "Default", + }, + "alerts": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-alerts"), + Title: "Test Alerts", + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Inject mock client + var sentURL string + ch.client = &mockTeamsClient{ + sendFunc: func(ctx context.Context, webhookURL string, message goteamsnotify.TeamsMessage) error { + sentURL = webhookURL + return nil + }, + } + + ctx := context.Background() + _ = ch.Start(ctx) + defer ch.Stop(ctx) + + msg := bus.OutboundMessage{Content: "Hello Teams!", ChatID: "alerts"} + + _, err = ch.Send(ctx, msg) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if sentURL != "https://example.com/webhook-alerts" { + t.Errorf("expected webhook URL 'https://example.com/webhook-alerts', got %q", sentURL) + } +} + +func TestTeamsWebhookChannel_SendError(t *testing.T) { + msgBus := bus.NewMessageBus() + ch, err := NewTeamsWebhookChannel(config.TeamsWebhookConfig{ + Enabled: true, + Webhooks: map[string]config.TeamsWebhookTarget{ + "default": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-default"), + }, + "alerts": { + WebhookURL: *config.NewSecureString("https://example.com/webhook-alerts"), + }, + }, + }, msgBus) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Inject mock client that returns an error + ch.client = &mockTeamsClient{ + sendFunc: func(ctx context.Context, webhookURL string, message goteamsnotify.TeamsMessage) error { + return errors.New("error on notification: 401 Unauthorized, forbidden") + }, + } + + ctx := context.Background() + _ = ch.Start(ctx) + defer ch.Stop(ctx) + + msg := bus.OutboundMessage{Content: "test", ChatID: "alerts"} + + _, err = ch.Send(ctx, msg) + if err == nil { + t.Error("expected error from failed send") + } +} + +func TestSplitContentWithTables(t *testing.T) { + tests := []struct { + name string + content string + wantSegs int + wantTbl int // number of table segments + }{ + { + name: "no tables", + content: "Just some text\nwith multiple lines", + wantSegs: 1, + wantTbl: 0, + }, + { + name: "single table", + content: `| Col1 | Col2 | +|------|------| +| A | B | +| C | D |`, + wantSegs: 1, + wantTbl: 1, + }, + { + name: "text before table", + content: `Here is some text. + +| Col1 | Col2 | +|------|------| +| A | B |`, + wantSegs: 2, + wantTbl: 1, + }, + { + name: "text before and after table", + content: `Before table. + +| Col1 | Col2 | +|------|------| +| A | B | + +After table.`, + wantSegs: 3, + wantTbl: 1, + }, + { + name: "multiple tables", + content: `First table: + +| A | B | +|---|---| +| 1 | 2 | + +Second table: + +| X | Y | +|---|---| +| 3 | 4 |`, + wantSegs: 4, + wantTbl: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + segs := splitContentWithTables(tt.content) + if len(segs) != tt.wantSegs { + t.Errorf("got %d segments, want %d", len(segs), tt.wantSegs) + } + tableCount := 0 + for _, s := range segs { + if s.isTable { + tableCount++ + } + } + if tableCount != tt.wantTbl { + t.Errorf("got %d tables, want %d", tableCount, tt.wantTbl) + } + }) + } +} + +func TestParseMarkdownTable(t *testing.T) { + tableStr := `| Name | Value | +|------|-------| +| foo | 123 | +| bar | 456 |` + + elem, err := parseMarkdownTable(tableStr) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if elem.Type != "Table" { + t.Errorf("expected type 'Table', got %q", elem.Type) + } + + // Should have 3 rows (header + 2 data rows) + if len(elem.Rows) != 3 { + t.Errorf("expected 3 rows, got %d", len(elem.Rows)) + } + + // Should have 2 columns with widths based on content length + if len(elem.Columns) != 2 { + t.Errorf("expected 2 columns, got %d", len(elem.Columns)) + } +} + +func TestParseMarkdownTableColumnWidths(t *testing.T) { + // Column widths are based on HEADER row only: + // Col1: "Description" (11 chars) + // Col2: "X" (1 char) + // Col3: "Amount" (6 chars) + tableStr := `| Description | X | Amount | +|-------------|---|--------| +| Short | Y | 100 | +| Longer text | Z | 50 |` + + elem, err := parseMarkdownTable(tableStr) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(elem.Columns) != 3 { + t.Fatalf("expected 3 columns, got %d", len(elem.Columns)) + } + + // Verify column widths are based on header content length + w1, ok1 := elem.Columns[0].Width.(int) + w2, ok2 := elem.Columns[1].Width.(int) + w3, ok3 := elem.Columns[2].Width.(int) + + if !ok1 || !ok2 || !ok3 { + t.Fatalf("expected int widths, got types: %T, %T, %T", + elem.Columns[0].Width, elem.Columns[1].Width, elem.Columns[2].Width) + } + + // Header lengths: "Description" = 11, "X" = 1, "Amount" = 6 + if w1 != 11 { + t.Errorf("expected col1 width 11 (from 'Description'), got %d", w1) + } + if w2 != 1 { + t.Errorf("expected col2 width 1 (from 'X'), got %d", w2) + } + if w3 != 6 { + t.Errorf("expected col3 width 6 (from 'Amount'), got %d", w3) + } +} + +func TestCalculateColumnWidths(t *testing.T) { + tests := []struct { + name string + maxLengths []int + wantWidths []int + }{ + { + name: "equal lengths", + maxLengths: []int{10, 10, 10}, + wantWidths: []int{10, 10, 10}, + }, + { + name: "varying lengths", + maxLengths: []int{5, 20, 10}, + wantWidths: []int{5, 20, 10}, + }, + { + name: "zero length gets minimum of 1", + maxLengths: []int{0, 5, 0}, + wantWidths: []int{1, 5, 1}, + }, + { + name: "empty input", + maxLengths: []int{}, + wantWidths: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cols := calculateColumnWidths(tt.maxLengths) + + if tt.wantWidths == nil { + if cols != nil { + t.Errorf("expected nil, got %v", cols) + } + return + } + + if len(cols) != len(tt.wantWidths) { + t.Fatalf("expected %d columns, got %d", len(tt.wantWidths), len(cols)) + } + + for i, col := range cols { + width, ok := col.Width.(int) + if !ok { + t.Errorf("column %d: expected int width, got %T", i, col.Width) + continue + } + if width != tt.wantWidths[i] { + t.Errorf("column %d: expected width %d, got %d", i, tt.wantWidths[i], width) + } + if col.Type != "TableColumnDefinition" { + t.Errorf("column %d: expected type 'TableColumnDefinition', got %q", i, col.Type) + } + } + }) + } +} + +func TestParseTableRow(t *testing.T) { + tests := []struct { + line string + want []string + }{ + {"| A | B | C |", []string{"A", "B", "C"}}, + {"|A|B|C|", []string{"A", "B", "C"}}, + {"| foo | bar |", []string{"foo", "bar"}}, + {"", nil}, + } + + for _, tt := range tests { + got := parseTableRow(tt.line) + if len(got) != len(tt.want) { + t.Errorf("parseTableRow(%q): got %v, want %v", tt.line, got, tt.want) + continue + } + for i := range got { + if got[i] != tt.want[i] { + t.Errorf("parseTableRow(%q)[%d]: got %q, want %q", tt.line, i, got[i], tt.want[i]) + } + } + } +} + +func TestIsSeparatorRow(t *testing.T) { + tests := []struct { + line string + want bool + }{ + {"|---|---|", true}, + {"| --- | --- |", true}, + {"|:---|---:|", true}, + {"| :---: | :---: |", true}, + {"| A | B |", false}, + {"| foo | bar |", false}, + } + + for _, tt := range tests { + got := isSeparatorRow(tt.line) + if got != tt.want { + t.Errorf("isSeparatorRow(%q): got %v, want %v", tt.line, got, tt.want) + } + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 7165246e5..fd4466b8c 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -24,20 +24,21 @@ var rrCounter atomic.Uint64 // CurrentVersion is the latest config schema version const CurrentVersion = 2 -// Config is the current config structure with version support +// Config is the current config structure with version support. type Config struct { - Version int `json:"version" yaml:"-"` // Config schema version for migration - Agents AgentsConfig `json:"agents" yaml:"-"` - Bindings []AgentBinding `json:"bindings,omitempty" yaml:"-"` - Session SessionConfig `json:"session,omitempty" yaml:"-"` - Channels ChannelsConfig `json:"channels" yaml:"channels"` - ModelList SecureModelList `json:"model_list" yaml:"model_list"` // New model-centric provider configuration - Gateway GatewayConfig `json:"gateway" yaml:"-"` - Hooks HooksConfig `json:"hooks,omitempty" yaml:"-"` - Tools ToolsConfig `json:"tools" yaml:",inline"` - Heartbeat HeartbeatConfig `json:"heartbeat" yaml:"-"` - Devices DevicesConfig `json:"devices" yaml:"-"` - Voice VoiceConfig `json:"voice" yaml:"-"` + Version int `json:"version" yaml:"-"` // Config schema version for migration + Isolation IsolationConfig `json:"isolation,omitempty" yaml:"-"` + Agents AgentsConfig `json:"agents" yaml:"-"` + Bindings []AgentBinding `json:"bindings,omitempty" yaml:"-"` + Session SessionConfig `json:"session,omitempty" yaml:"-"` + Channels ChannelsConfig `json:"channels" yaml:"channels"` + ModelList SecureModelList `json:"model_list" yaml:"model_list"` // New model-centric provider configuration + Gateway GatewayConfig `json:"gateway" yaml:"-"` + Hooks HooksConfig `json:"hooks,omitempty" yaml:"-"` + Tools ToolsConfig `json:"tools" yaml:",inline"` + Heartbeat HeartbeatConfig `json:"heartbeat" yaml:"-"` + Devices DevicesConfig `json:"devices" yaml:"-"` + Voice VoiceConfig `json:"voice" yaml:"-"` // BuildInfo contains build-time version information BuildInfo BuildInfo `json:"build_info,omitempty" yaml:"-"` @@ -45,6 +46,21 @@ type Config struct { sensitiveCache *SensitiveDataCache } +// IsolationConfig controls subprocess isolation for commands started by PicoClaw. +// It is applied by the isolation package rather than by sandboxing the main process. +type IsolationConfig struct { + Enabled bool `json:"enabled,omitempty"` + ExposePaths []ExposePath `json:"expose_paths,omitempty"` +} + +// ExposePath describes a host path that should remain visible inside the isolated +// child-process environment. This is currently implemented on Linux only. +type ExposePath struct { + Source string `json:"source"` + Target string `json:"target,omitempty"` + Mode string `json:"mode"` +} + // FilterSensitiveData filters sensitive values from content before sending to LLM. // This prevents the LLM from seeing its own credentials. // Uses strings.Replacer for O(n+m) performance (computed once per SecurityConfig). @@ -280,23 +296,24 @@ func (d *AgentDefaults) GetModelName() string { } type ChannelsConfig struct { - WhatsApp WhatsAppConfig `json:"whatsapp" yaml:"-"` - Telegram TelegramConfig `json:"telegram" yaml:"telegram,omitempty"` - Feishu FeishuConfig `json:"feishu" yaml:"feishu,omitempty"` - Discord DiscordConfig `json:"discord" yaml:"discord,omitempty"` - MaixCam MaixCamConfig `json:"maixcam" yaml:"-"` - QQ QQConfig `json:"qq" yaml:"qq,omitempty"` - DingTalk DingTalkConfig `json:"dingtalk" yaml:"dingtalk,omitempty"` - Slack SlackConfig `json:"slack" yaml:"slack,omitempty"` - Matrix MatrixConfig `json:"matrix" yaml:"matrix,omitempty"` - LINE LINEConfig `json:"line" yaml:"line,omitempty"` - OneBot OneBotConfig `json:"onebot" yaml:"onebot,omitempty"` - WeCom WeComConfig `json:"wecom" yaml:"wecom,omitempty" envPrefix:"PICOCLAW_CHANNELS_WECOM_"` - Weixin WeixinConfig `json:"weixin" yaml:"weixin,omitempty"` - Pico PicoConfig `json:"pico" yaml:"pico,omitempty"` - PicoClient PicoClientConfig `json:"pico_client" yaml:"pico_client,omitempty"` - IRC IRCConfig `json:"irc" yaml:"irc,omitempty"` - VK VKConfig `json:"vk" yaml:"vk,omitempty"` + WhatsApp WhatsAppConfig `json:"whatsapp" yaml:"-"` + Telegram TelegramConfig `json:"telegram" yaml:"telegram,omitempty"` + Feishu FeishuConfig `json:"feishu" yaml:"feishu,omitempty"` + Discord DiscordConfig `json:"discord" yaml:"discord,omitempty"` + MaixCam MaixCamConfig `json:"maixcam" yaml:"-"` + QQ QQConfig `json:"qq" yaml:"qq,omitempty"` + DingTalk DingTalkConfig `json:"dingtalk" yaml:"dingtalk,omitempty"` + Slack SlackConfig `json:"slack" yaml:"slack,omitempty"` + Matrix MatrixConfig `json:"matrix" yaml:"matrix,omitempty"` + LINE LINEConfig `json:"line" yaml:"line,omitempty"` + OneBot OneBotConfig `json:"onebot" yaml:"onebot,omitempty"` + WeCom WeComConfig `json:"wecom" yaml:"wecom,omitempty" envPrefix:"PICOCLAW_CHANNELS_WECOM_"` + Weixin WeixinConfig `json:"weixin" yaml:"weixin,omitempty"` + Pico PicoConfig `json:"pico" yaml:"pico,omitempty"` + PicoClient PicoClientConfig `json:"pico_client" yaml:"pico_client,omitempty"` + IRC IRCConfig `json:"irc" yaml:"irc,omitempty"` + VK VKConfig `json:"vk" yaml:"vk,omitempty"` + TeamsWebhook TeamsWebhookConfig `json:"teams_webhook" yaml:"teams_webhook,omitempty"` } // GroupTriggerConfig controls when the bot responds in group chats. @@ -566,6 +583,19 @@ func (c *VKConfig) SetToken(token string) { c.Token = *NewSecureString(token) } +// TeamsWebhookConfig configures the output-only Microsoft Teams webhook channel. +// Multiple webhook targets can be configured and selected via ChatID at send time. +type TeamsWebhookConfig struct { + Enabled bool `json:"enabled" yaml:"-" env:"PICOCLAW_CHANNELS_TEAMS_WEBHOOK_ENABLED"` + Webhooks map[string]TeamsWebhookTarget `json:"webhooks" yaml:"webhooks,omitempty"` +} + +// TeamsWebhookTarget represents a single Teams webhook destination. +type TeamsWebhookTarget struct { + WebhookURL SecureString `json:"webhook_url,omitzero" yaml:"webhook_url,omitempty"` + Title string `json:"title,omitempty" yaml:"-"` +} + type HeartbeatConfig struct { Enabled bool `json:"enabled" env:"PICOCLAW_HEARTBEAT_ENABLED"` Interval int `json:"interval" env:"PICOCLAW_HEARTBEAT_INTERVAL"` // minutes, min 5 @@ -605,11 +635,12 @@ type ModelConfig struct { Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers // Optional optimizations - RPM int `json:"rpm,omitempty"` // Requests per minute limit - MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive - ExtraBody map[string]any `json:"extra_body,omitempty"` // Additional fields to inject into request body + RPM int `json:"rpm,omitempty"` // Requests per minute limit + MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive + ExtraBody map[string]any `json:"extra_body,omitempty"` // Additional fields to inject into request body + CustomHeaders map[string]string `json:"custom_headers,omitempty"` // Additional headers to inject into every HTTP request APIKeys SecureStrings `json:"api_keys,omitzero" yaml:"api_keys,omitempty"` // API authentication keys (multiple keys for failover) @@ -1279,6 +1310,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ExtraBody: m.ExtraBody, + CustomHeaders: m.CustomHeaders, isVirtual: true, } expanded = append(expanded, additionalEntry) @@ -1299,6 +1331,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ExtraBody: m.ExtraBody, + CustomHeaders: m.CustomHeaders, APIKeys: SimpleSecureStrings(keys[0]), } diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 8e58a684e..f0449d98f 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -852,6 +852,37 @@ func TestDefaultConfig_WorkspacePath_WithPicoclawHome(t *testing.T) { } } +func TestDefaultConfig_IsolationEnabled(t *testing.T) { + cfg := DefaultConfig() + if cfg.Isolation.Enabled { + t.Fatal("DefaultConfig().Isolation.Enabled should be false") + } +} + +func TestConfig_UnmarshalIsolation(t *testing.T) { + cfg := DefaultConfig() + raw := []byte(`{ + "isolation": { + "enabled": false, + "expose_paths": [ + {"source":"/src","target":"/dst","mode":"ro"} + ] + } + }`) + if err := json.Unmarshal(raw, cfg); err != nil { + t.Fatalf("json.Unmarshal isolation config: %v", err) + } + if cfg.Isolation.Enabled { + t.Fatal("Isolation.Enabled should be false after unmarshal") + } + if len(cfg.Isolation.ExposePaths) != 1 { + t.Fatalf("ExposePaths len = %d, want 1", len(cfg.Isolation.ExposePaths)) + } + if got := cfg.Isolation.ExposePaths[0]; got.Source != "/src" || got.Target != "/dst" || got.Mode != "ro" { + t.Fatalf("ExposePaths[0] = %+v, want source=/src target=/dst mode=ro", got) + } +} + // TestFlexibleStringSlice_UnmarshalText tests UnmarshalText with various comma separators func TestFlexibleStringSlice_UnmarshalText(t *testing.T) { tests := []struct { @@ -1528,6 +1559,42 @@ func TestModelConfig_ExtraBodyRoundTrip(t *testing.T) { } } +func TestModelConfig_CustomHeadersRoundTrip(t *testing.T) { + dir := t.TempDir() + cfgPath := filepath.Join(dir, "config.json") + + cfg := &Config{ + Version: CurrentVersion, + ModelList: []*ModelConfig{ + { + ModelName: "test-model", + Model: "openai/test", + APIKeys: SimpleSecureStrings("sk-test"), + CustomHeaders: map[string]string{"X-Source": "coding-plan", "X-Agent": "openclaw"}, + }, + }, + } + + if err := SaveConfig(cfgPath, cfg); err != nil { + t.Fatalf("SaveConfig error: %v", err) + } + + loaded, err := LoadConfig(cfgPath) + if err != nil { + t.Fatalf("LoadConfig error: %v", err) + } + + if loaded.ModelList[0].CustomHeaders == nil { + t.Fatal("CustomHeaders should not be nil after round-trip") + } + if got := loaded.ModelList[0].CustomHeaders["X-Source"]; got != "coding-plan" { + t.Errorf("CustomHeaders[X-Source] = %q, want coding-plan", got) + } + if got := loaded.ModelList[0].CustomHeaders["X-Agent"]; got != "openclaw" { + t.Errorf("CustomHeaders[X-Agent] = %q, want openclaw", got) + } +} + func TestDefaultConfig_MinimaxExtraBody(t *testing.T) { cfg := DefaultConfig() diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index c2e1a31f3..bb073d436 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -17,6 +17,11 @@ func DefaultConfig() *Config { return &Config{ Version: CurrentVersion, + // Isolation is opt-in so existing installations keep their current behavior + // until the user explicitly enables subprocess sandboxing. + Isolation: IsolationConfig{ + Enabled: false, + }, Agents: AgentsConfig{ Defaults: AgentDefaults{ Workspace: workspacePath, diff --git a/pkg/gateway/gateway.go b/pkg/gateway/gateway.go index 509b5d37e..8be84bdf6 100644 --- a/pkg/gateway/gateway.go +++ b/pkg/gateway/gateway.go @@ -28,6 +28,7 @@ import ( "github.com/sipeed/picoclaw/pkg/channels/pico" _ "github.com/sipeed/picoclaw/pkg/channels/qq" _ "github.com/sipeed/picoclaw/pkg/channels/slack" + _ "github.com/sipeed/picoclaw/pkg/channels/teams_webhook" _ "github.com/sipeed/picoclaw/pkg/channels/telegram" _ "github.com/sipeed/picoclaw/pkg/channels/vk" _ "github.com/sipeed/picoclaw/pkg/channels/wecom" diff --git a/pkg/health/server.go b/pkg/health/server.go index 2602cb965..a152d8ab1 100644 --- a/pkg/health/server.go +++ b/pkg/health/server.go @@ -7,6 +7,7 @@ import ( "fmt" "maps" "net/http" + "os" "sync" "time" ) @@ -31,6 +32,7 @@ type Check struct { type StatusResponse struct { Status string `json:"status"` Uptime string `json:"uptime"` + PID int `json:"pid,omitempty"` Checks map[string]Check `json:"checks,omitempty"` } @@ -170,6 +172,7 @@ func (s *Server) healthHandler(w http.ResponseWriter, r *http.Request) { resp := StatusResponse{ Status: "ok", Uptime: uptime.String(), + PID: os.Getpid(), } json.NewEncoder(w).Encode(resp) diff --git a/pkg/isolation/README.md b/pkg/isolation/README.md new file mode 100644 index 000000000..de16ce505 --- /dev/null +++ b/pkg/isolation/README.md @@ -0,0 +1,238 @@ +# `pkg/isolation` + +`pkg/isolation` provides process-level isolation for child processes started by `picoclaw`. + +It does not sandbox the main `picoclaw` process itself. + +## Scope + +The current scope is the child-process startup path: + +- `exec` tool +- CLI providers such as `claude-cli` and `codex-cli` +- process hooks +- MCP `stdio` servers + +## One-Sentence Model + +- The `picoclaw` main process still runs in the host environment. +- Every child process should enter the shared `pkg/isolation` startup path first. +- The startup path applies platform-specific isolation according to config. + +## Architecture + +The implementation has four layers: + +1. Configuration layer: reads `config.Config.Isolation` and injects it through `isolation.Configure(cfg)`. +2. Instance layout layer: resolves `config.GetHome()`, prepares instance directories, and builds the runtime user environment. +3. Platform backend layer: Linux uses `bwrap`; Windows uses a restricted token, low integrity, and a `Job Object`; other platforms are not implemented. +4. Unified startup layer: `PrepareCommand(cmd)`, `Start(cmd)`, and `Run(cmd)`. + +All integrations that spawn subprocesses should reuse these helpers instead of calling `cmd.Start` or `cmd.Run` directly. + +## Configuration + +Isolation lives under: + +```json +{ + "isolation": { + "enabled": false, + "expose_paths": [] + } +} +``` + +Field meanings: + +- `enabled`: enables or disables subprocess isolation. Default: `false`. +- `expose_paths`: explicitly exposes host paths inside the isolated environment. It only matters when `enabled=true`. This is currently supported on Linux only. + +Example: + +```json +{ + "isolation": { + "enabled": true, + "expose_paths": [ + { + "source": "/opt/toolchains/go", + "target": "/opt/toolchains/go", + "mode": "ro" + }, + { + "source": "/data/shared-assets", + "target": "/opt/picoclaw-instance-a/workspace/assets", + "mode": "rw" + } + ] + } +} +``` + +Rules for `expose_paths`: + +- `source` is a host path. +- `target` is the path inside the isolated environment. +- `mode` must be `ro` or `rw`. +- When `target` is empty, it defaults to `source`. +- Only one final rule may exist for the same `target`. +- Later-loaded config overrides earlier rules for the same `target`. + +Platform note: + +- Linux uses a real `source -> target` mount view. +- Windows does not currently support `expose_paths`. + +## Instance Root And Directories + +The instance root follows `config.GetHome()`: + +- If `PICOCLAW_HOME` is set, use it. +- Otherwise use the default `.picoclaw` directory under the user home. + +If `config.GetHome()` falls back to `.` while isolation is enabled, startup should fail. + +Default instance directories include: + +- instance root +- `skills` +- `logs` +- `cache` +- `state` +- `runtime-user-env` + +`workspace` is derived from `cfg.WorkspacePath()` when configured, otherwise from the default workspace rule. + +Windows also prepares: + +- `runtime-user-env/AppData/Roaming` +- `runtime-user-env/AppData/Local` + +## User Environment Redirect + +When isolation is enabled, child processes receive a redirected per-instance user environment. + +Linux variables: + +- `HOME` +- `TMPDIR` +- `XDG_CONFIG_HOME` +- `XDG_CACHE_HOME` +- `XDG_STATE_HOME` + +Windows variables: + +- `USERPROFILE` +- `HOME` +- `TEMP` +- `TMP` +- `APPDATA` +- `LOCALAPPDATA` + +These paths point into `runtime-user-env` under the instance root. + +## Platform Behavior + +### Linux + +The Linux backend currently depends on `bwrap` (`bubblewrap`). + +Capabilities: + +- minimal filesystem view +- `ipc` namespace isolation +- redirected child-process user environment +- `source -> target` read-only or read-write mounts + +Default mounts include the instance root plus the minimum runtime system paths such as `/usr`, `/bin`, `/lib`, `/lib64`, and `/etc/resolv.conf`. + +At runtime, PicoClaw also adds the executable path, its directory, the effective working directory, and absolute path arguments when needed. + +There is no automatic fallback when `bwrap` is missing. + +Install examples: + +- `apt install bubblewrap` +- `dnf install bubblewrap` +- `yum install bubblewrap` +- `pacman -S bubblewrap` +- `apk add bubblewrap` + +If isolation must be disabled temporarily: + +```json +{ + "isolation": { + "enabled": false + } +} +``` + +Disabling isolation increases the risk that child processes can access or modify more host files. + +### Windows + +Windows isolation currently supports process-level restrictions such as restricted tokens, low integrity, job objects, and redirected user-environment directories. + +`expose_paths` is not currently supported on Windows. If it is configured, startup should fail instead of pretending the paths were exposed. + +The Windows backend currently uses: + +- a restricted primary token +- low integrity level +- a `Job Object` +- redirected child-process user environment + +It does not currently implement true `source -> target` filesystem remapping. + +### macOS And Other Platforms + +They are not implemented yet. + +When isolation is explicitly enabled on an unsupported platform, the higher-level runtime should surface that as an unsupported configuration instead of pretending isolation succeeded. + +## Logging And Debugging + +When isolation is enabled, PicoClaw logs the generated isolation plan. + +Linux log name: + +- `linux isolation mount plan` + +Windows log name: + +- `windows isolation access rules` + +If you suspect isolation is ineffective, check whether unexpected host paths appear in those logs. + +## Relationship To `restrict_to_workspace` + +- `restrict_to_workspace` limits the paths an agent is normally allowed to access. +- `pkg/isolation` limits what a child process can see and where its user environment points. + +They complement each other and do not replace each other. + +## Current Limits + +- Linux isolation is implemented with `bwrap`, not a custom in-process isolation runtime. +- Linux does not currently enable a dedicated `pid` namespace by default. +- Windows does not yet implement full host ACL enforcement for every allowed or denied path. +- macOS is not implemented. +- The current design isolates child processes, not the main `picoclaw` process. + +## Suggested Reading Order + +If you are new to this code, read it in this order: + +1. `pkg/config/config.go` +2. `pkg/isolation/runtime.go` +3. `pkg/isolation/platform_linux.go` +4. `pkg/isolation/platform_windows.go` +5. Call sites: +6. `pkg/tools/shell.go` +7. `pkg/providers/*.go` +8. `pkg/agent/hook_process.go` +9. `pkg/mcp/manager.go` + +That path gives the fastest overview of the configuration model, runtime flow, and platform-specific limits. diff --git a/pkg/isolation/README_CN.md b/pkg/isolation/README_CN.md new file mode 100644 index 000000000..0529a84bd --- /dev/null +++ b/pkg/isolation/README_CN.md @@ -0,0 +1,238 @@ +# `pkg/isolation` + +`pkg/isolation` 为 `picoclaw` 启动的子进程提供进程级隔离能力。 + +它当前不会把 `picoclaw` 主进程自身放进沙箱中运行。 + +## 生效范围 + +当前生效范围是子进程启动链路： + +- `exec` 工具 +- `claude-cli`、`codex-cli` 等 CLI provider +- 进程型 hooks +- MCP `stdio` server + +## 一句话理解 + +- `picoclaw` 主进程仍运行在宿主环境中。 +- 所有子进程都应先经过 `pkg/isolation` 的统一启动入口。 +- 入口会根据配置和平台，为子进程施加对应隔离。 + +## 架构 + +当前实现可以分为四层： + +1. 配置层：读取 `config.Config.Isolation`，并通过 `isolation.Configure(cfg)` 注入运行时。 +2. 实例目录层：解析 `config.GetHome()`，准备实例目录，并构建运行时用户环境目录。 +3. 平台后端层：Linux 使用 `bwrap`；Windows 使用受限 token、低完整性级别和 `Job Object`；其他平台未实现。 +4. 统一启动层：`PrepareCommand(cmd)`、`Start(cmd)`、`Run(cmd)`。 + +所有启动子进程的接入点都应复用这组入口，而不是各自直接调用 `cmd.Start` 或 `cmd.Run`。 + +## 配置 + +隔离配置位于： + +```json +{ + "isolation": { + "enabled": false, + "expose_paths": [] + } +} +``` + +字段说明： + +- `enabled`：是否启用子进程隔离。默认值：`false`。 +- `expose_paths`：显式把宿主路径带入隔离环境。仅在 `enabled=true` 时生效。目前只在 Linux 上支持。 + +示例： + +```json +{ + "isolation": { + "enabled": true, + "expose_paths": [ + { + "source": "/opt/toolchains/go", + "target": "/opt/toolchains/go", + "mode": "ro" + }, + { + "source": "/data/shared-assets", + "target": "/opt/picoclaw-instance-a/workspace/assets", + "mode": "rw" + } + ] + } +} +``` + +`expose_paths` 规则： + +- `source`：宿主机路径。 +- `target`：隔离环境内的目标路径。 +- `mode`：只能是 `ro` 或 `rw`。 +- `target` 为空时，默认等于 `source`。 +- 同一个 `target` 最终只能保留一条规则。 +- 后加载的配置会覆盖先加载的同目标规则。 + +平台说明： + +- Linux 会真实使用 `source -> target` 挂载视图。 +- Windows 当前不支持 `expose_paths`。 + +## 实例根与目录 + +实例根遵循 `config.GetHome()`： + +- 如果设置了 `PICOCLAW_HOME`，使用该值。 +- 否则默认使用用户目录下的 `.picoclaw`。 + +如果 `config.GetHome()` 在隔离开启时最终回退到当前目录 `.`，启动应直接失败。 + +默认实例目录包括： + +- 实例根本身 +- `skills` +- `logs` +- `cache` +- `state` +- `runtime-user-env` + +`workspace` 优先使用 `cfg.WorkspacePath()` 的结果；未显式配置时才按默认规则派生。 + +Windows 还会额外准备： + +- `runtime-user-env/AppData/Roaming` +- `runtime-user-env/AppData/Local` + +## 用户环境重定向 + +隔离开启后，子进程会收到重定向到实例目录下的独立用户环境。 + +Linux 注入变量： + +- `HOME` +- `TMPDIR` +- `XDG_CONFIG_HOME` +- `XDG_CACHE_HOME` +- `XDG_STATE_HOME` + +Windows 注入变量： + +- `USERPROFILE` +- `HOME` +- `TEMP` +- `TMP` +- `APPDATA` +- `LOCALAPPDATA` + +这些路径都会指向实例根下的 `runtime-user-env`。 + +## 平台行为 + +### Linux + +Linux 后端当前依赖 `bwrap`（`bubblewrap`）。 + +能力： + +- 最小文件系统视图 +- `ipc namespace` +- 子进程用户环境重定向 +- `source -> target` 只读或读写挂载 + +默认映射包括实例根，以及 `/usr`、`/bin`、`/lib`、`/lib64`、`/etc/resolv.conf` 等最小运行时系统路径。 + +运行时还会按需补充可执行文件本身、其所在目录、生效后的工作目录，以及命令行中的绝对路径参数。 + +缺少 `bwrap` 时不会自动回退。 + +安装示例： + +- `apt install bubblewrap` +- `dnf install bubblewrap` +- `yum install bubblewrap` +- `pacman -S bubblewrap` +- `apk add bubblewrap` + +如果需要临时关闭隔离： + +```json +{ + "isolation": { + "enabled": false + } +} +``` + +关闭隔离后，子进程访问或修改更多宿主文件的风险会明显上升。 + +### Windows + +Windows 隔离当前提供的是进程级限制，例如 restricted token、low integrity、job object，以及用户环境目录重定向。 + +`expose_paths` 目前不支持 Windows。如果配置了该字段，启动应直接失败，而不是假装这些路径已经被暴露进隔离环境。 + +Windows 后端当前使用： + +- 受限 primary token +- 低完整性级别 +- `Job Object` +- 子进程用户环境重定向 + +它当前不会实现真正的 `source -> target` 文件系统重映射。 + +### macOS 与其他平台 + +当前尚未实现。 + +当在未支持的平台上显式开启隔离时，上层运行时应将其视为不支持的配置，而不是假装隔离成功。 + +## 日志与排障 + +隔离开启后，PicoClaw 会打印生成后的隔离计划，便于排障。 + +Linux 日志名： + +- `linux isolation mount plan` + +Windows 日志名： + +- `windows isolation access rules` + +如果你怀疑隔离未生效，先检查这些日志里是否出现了不应暴露的宿主路径。 + +## 与 `restrict_to_workspace` 的关系 + +- `restrict_to_workspace` 限制的是 agent 默认可访问的路径。 +- `pkg/isolation` 限制的是子进程运行时能看到什么文件系统，以及它的用户环境指向哪里。 + +两者互补，不互相替代。 + +## 当前限制 + +- Linux 基于 `bwrap` 实现，而不是纯内建 isolation runtime。 +- Linux 当前没有默认启用独立的 `pid namespace`。 +- Windows 还没有对所有允许/拒绝路径做完整 ACL 落地。 +- macOS 尚未实现。 +- 当前隔离的是子进程，不是 `picoclaw` 主进程自身。 + +## 建议阅读顺序 + +如果你是第一次看这部分代码，建议按这个顺序阅读： + +1. `pkg/config/config.go` +2. `pkg/isolation/runtime.go` +3. `pkg/isolation/platform_linux.go` +4. `pkg/isolation/platform_windows.go` +5. 调用点： +6. `pkg/tools/shell.go` +7. `pkg/providers/*.go` +8. `pkg/agent/hook_process.go` +9. `pkg/mcp/manager.go` + +这样能最快建立对配置模型、运行流程和平台边界的整体理解。 diff --git a/pkg/isolation/platform_linux.go b/pkg/isolation/platform_linux.go new file mode 100644 index 000000000..9a282a4ad --- /dev/null +++ b/pkg/isolation/platform_linux.go @@ -0,0 +1,264 @@ +//go:build linux + +package isolation + +import ( + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +func applyPlatformIsolation(cmd *exec.Cmd, isolation config.IsolationConfig, root string) error { + if !isolation.Enabled { + return nil + } + // Bubblewrap is the only supported Linux backend right now. Fail closed when + // it is unavailable instead of silently running the child process unisolated. + bwrapPath, err := exec.LookPath("bwrap") + if err != nil { + hint := bwrapInstallHint() + disableHint := `set "isolation.enabled": false in config.json` + logger.WarnCF("isolation", "bubblewrap is required for Linux isolation", + map[string]any{ + "binary": "bwrap", + "install": hint, + "disable_isolation": disableHint, + "risk": "disabling isolation lets child processes run without Linux filesystem isolation", + }) + return fmt.Errorf( + "linux isolation requires bwrap and does not fall back automatically: %w; install bubblewrap with one of: %s; or disable isolation by setting %s; disabling isolation means child processes can run without Linux filesystem isolation and may access or modify more host files", + err, + hint, + disableHint, + ) + } + if cmd == nil || cmd.Path == "" || len(cmd.Args) == 0 { + return nil + } + + originalPath := cmd.Path + originalArgs := append([]string{}, cmd.Args...) + _, execDir, err := resolveLinuxWorkingDir(cmd.Dir, originalPath) + if err != nil { + return err + } + resolvedPath, err := resolveLinuxCommandPath(originalPath, execDir) + if err != nil { + return err + } + + // Start from the configured mount plan, then add only the executable, its + // resolved path, the effective working directory, and any absolute path + // arguments needed to preserve the original command semantics. + plan := BuildLinuxMountPlan(root, isolation.ExposePaths) + plan = ensureLinuxMountRule(plan, resolvedPath, resolvedPath, "ro") + plan = ensureLinuxMountRule(plan, filepath.Dir(resolvedPath), filepath.Dir(resolvedPath), "ro") + if resolved, resolveErr := filepath.EvalSymlinks(resolvedPath); resolveErr == nil && resolved != resolvedPath { + plan = ensureLinuxMountRule(plan, resolved, resolved, "ro") + plan = ensureLinuxMountRule(plan, filepath.Dir(resolved), filepath.Dir(resolved), "ro") + } + if execDir != "" { + plan = ensureLinuxMountRule(plan, execDir, execDir, "rw") + if resolved, resolveErr := filepath.EvalSymlinks(execDir); resolveErr == nil && resolved != execDir { + plan = ensureLinuxMountRule(plan, resolved, resolved, "rw") + } + } + plan = appendLinuxArgumentMounts(plan, originalArgs[1:]) + logger.DebugCF("isolation", "linux isolation mount plan", + map[string]any{ + "root": root, + "command": resolvedPath, + "working_dir": execDir, + "mounts": formatLinuxMountPlan(plan), + }) + bwrapArgs, err := buildLinuxBwrapArgs(originalPath, resolvedPath, originalArgs, execDir, plan) + if err != nil { + return err + } + + cmd.Path = bwrapPath + cmd.Args = bwrapArgs + cmd.Dir = "" + return nil +} + +func bwrapInstallHint() string { + return "apt install bubblewrap; dnf install bubblewrap; yum install bubblewrap; pacman -S bubblewrap; apk add bubblewrap" +} + +// formatLinuxMountPlan reshapes the internal plan for structured logging. +func formatLinuxMountPlan(plan []MountRule) []map[string]string { + formatted := make([]map[string]string, 0, len(plan)) + for _, rule := range plan { + formatted = append(formatted, map[string]string{ + "source": rule.Source, + "target": rule.Target, + "mode": rule.Mode, + }) + } + return formatted +} + +func postStartPlatformIsolation(cmd *exec.Cmd, isolation config.IsolationConfig, root string) error { + return nil +} + +func cleanupPendingPlatformResources(cmd *exec.Cmd) { +} + +// buildLinuxBwrapArgs translates the mount plan into the bubblewrap command +// line that re-executes the original process inside the isolated mount view. +func buildLinuxBwrapArgs( + originalPath string, + resolvedPath string, + originalArgs []string, + execDir string, + plan []MountRule, +) ([]string, error) { + bwrapArgs := []string{ + "bwrap", + "--die-with-parent", + "--unshare-ipc", + "--proc", "/proc", + "--dev", "/dev", + } + for _, rule := range plan { + flag, err := linuxBindFlag(rule) + if err != nil { + return nil, err + } + bwrapArgs = append(bwrapArgs, flag, rule.Source, rule.Target) + } + if execDir != "" { + bwrapArgs = append(bwrapArgs, "--chdir", execDir) + } + execPath := originalPath + if isRelativeCommandPath(originalPath) { + execPath = resolvedPath + } + bwrapArgs = append(bwrapArgs, "--", execPath) + if len(originalArgs) > 1 { + bwrapArgs = append(bwrapArgs, originalArgs[1:]...) + } + return bwrapArgs, nil +} + +func resolveLinuxWorkingDir(originalDir, originalPath string) (string, string, error) { + if originalDir != "" { + resolved, err := filepath.Abs(originalDir) + if err != nil { + return "", "", fmt.Errorf("resolve command dir %s: %w", originalDir, err) + } + return resolved, resolved, nil + } + if !isRelativeCommandPath(originalPath) { + return "", "", nil + } + wd, err := os.Getwd() + if err != nil { + return "", "", fmt.Errorf("resolve current working dir: %w", err) + } + return "", wd, nil +} + +func resolveLinuxCommandPath(originalPath, execDir string) (string, error) { + if filepath.IsAbs(originalPath) || !isRelativeCommandPath(originalPath) { + return filepath.Clean(originalPath), nil + } + base := execDir + if base == "" { + var err error + base, err = os.Getwd() + if err != nil { + return "", fmt.Errorf("resolve current working dir: %w", err) + } + } + return filepath.Clean(filepath.Join(base, originalPath)), nil +} + +func appendLinuxArgumentMounts(plan []MountRule, args []string) []MountRule { + for _, arg := range args { + path, ok := linuxArgumentPath(arg) + if !ok { + continue + } + clean := filepath.Clean(path) + if info, err := os.Stat(clean); err == nil { + mode := "ro" + if info.IsDir() { + mode = "rw" + } + plan = ensureLinuxMountRule(plan, clean, clean, mode) + if resolved, resolveErr := filepath.EvalSymlinks(clean); resolveErr == nil && resolved != clean { + plan = ensureLinuxMountRule(plan, resolved, resolved, mode) + } + continue + } else if !errors.Is(err, os.ErrNotExist) { + continue + } + parent := filepath.Dir(clean) + if parent == clean { + continue + } + if _, err := os.Stat(parent); err == nil { + plan = ensureLinuxMountRule(plan, parent, parent, "rw") + } + } + return plan +} + +func linuxArgumentPath(arg string) (string, bool) { + if filepath.IsAbs(arg) { + return arg, true + } + idx := strings.IndexRune(arg, '=') + if idx <= 0 || idx == len(arg)-1 { + return "", false + } + value := arg[idx+1:] + if !filepath.IsAbs(value) { + return "", false + } + return value, true +} + +func isRelativeCommandPath(path string) bool { + return !filepath.IsAbs(path) && strings.ContainsRune(path, filepath.Separator) +} + +// ensureLinuxMountRule appends a mount rule unless another rule already owns +// the same target path. +func ensureLinuxMountRule(plan []MountRule, source, target, mode string) []MountRule { + cleanSource := filepath.Clean(source) + cleanTarget := filepath.Clean(target) + for _, rule := range plan { + if filepath.Clean(rule.Target) == cleanTarget { + return plan + } + } + return append(plan, MountRule{Source: cleanSource, Target: cleanTarget, Mode: mode}) +} + +// linuxBindFlag selects the correct bubblewrap bind flag based on mount mode. +func linuxBindFlag(rule MountRule) (string, error) { + info, err := os.Stat(rule.Source) + if err != nil { + return "", fmt.Errorf("stat linux mount source %s: %w", rule.Source, err) + } + if !info.IsDir() { + if rule.Mode == "rw" { + return "--bind", nil + } + return "--ro-bind", nil + } + if rule.Mode == "rw" { + return "--bind", nil + } + return "--ro-bind", nil +} diff --git a/pkg/isolation/platform_linux_test.go b/pkg/isolation/platform_linux_test.go new file mode 100644 index 000000000..2dcca96ce --- /dev/null +++ b/pkg/isolation/platform_linux_test.go @@ -0,0 +1,148 @@ +//go:build linux + +package isolation + +import ( + "os" + "path/filepath" + "testing" + + "github.com/sipeed/picoclaw/pkg/config" +) + +func TestBuildLinuxBwrapArgs_IncludesNamespaceFlagsAndExec(t *testing.T) { + root := t.TempDir() + binaryDir := filepath.Join(root, "bin") + if err := os.MkdirAll(binaryDir, 0o755); err != nil { + t.Fatal(err) + } + binaryPath := filepath.Join(binaryDir, "tool") + if err := os.WriteFile(binaryPath, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil { + t.Fatal(err) + } + plan := BuildLinuxMountPlan(root, []config.ExposePath{{Source: binaryDir, Target: binaryDir, Mode: "ro"}}) + args, err := buildLinuxBwrapArgs(binaryPath, binaryPath, []string{binaryPath, "--flag"}, root, plan) + if err != nil { + t.Fatalf("buildLinuxBwrapArgs() error = %v", err) + } + hasNet := false + hasIPC := false + hasExec := false + for i := range args { + switch args[i] { + case "--unshare-net": + hasNet = true + case "--unshare-ipc": + hasIPC = true + case "--": + if i+1 < len(args) && args[i+1] == binaryPath { + hasExec = true + } + } + } + if hasNet { + t.Fatalf("bwrap args should not unshare net by default: %v", args) + } + if !hasIPC || !hasExec { + t.Fatalf("bwrap args missing required items: %v", args) + } +} + +func TestResolveLinuxWorkingDir_ResolvesRelativeDir(t *testing.T) { + cwd := t.TempDir() + previous, err := os.Getwd() + if err != nil { + t.Fatal(err) + } + defer func() { + if chdirErr := os.Chdir(previous); chdirErr != nil { + t.Fatalf("restore cwd: %v", chdirErr) + } + }() + if chdirErr := os.Chdir(cwd); chdirErr != nil { + t.Fatal(chdirErr) + } + + resolvedDir, execDir, err := resolveLinuxWorkingDir("./hooks", "./hook.sh") + if err != nil { + t.Fatalf("resolveLinuxWorkingDir() error = %v", err) + } + want := filepath.Join(cwd, "hooks") + if resolvedDir != want || execDir != want { + t.Fatalf("resolveLinuxWorkingDir() = (%q, %q), want (%q, %q)", resolvedDir, execDir, want, want) + } +} + +func TestResolveLinuxCommandPath_UsesExecDirForRelativeCommand(t *testing.T) { + execDir := filepath.Join(t.TempDir(), "hooks") + got, err := resolveLinuxCommandPath("./hook.sh", execDir) + if err != nil { + t.Fatalf("resolveLinuxCommandPath() error = %v", err) + } + want := filepath.Join(execDir, "hook.sh") + if got != want { + t.Fatalf("resolveLinuxCommandPath() = %q, want %q", got, want) + } +} + +func TestBuildLinuxBwrapArgs_UsesResolvedPathForRelativeCommand(t *testing.T) { + root := t.TempDir() + execDir := filepath.Join(root, "hooks") + if err := os.MkdirAll(execDir, 0o755); err != nil { + t.Fatal(err) + } + resolvedPath := filepath.Join(execDir, "hook.sh") + if err := os.WriteFile(resolvedPath, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil { + t.Fatal(err) + } + plan := []MountRule{ + {Source: execDir, Target: execDir, Mode: "rw"}, + {Source: resolvedPath, Target: resolvedPath, Mode: "ro"}, + } + args, err := buildLinuxBwrapArgs("./hook.sh", resolvedPath, []string{"./hook.sh"}, execDir, plan) + if err != nil { + t.Fatalf("buildLinuxBwrapArgs() error = %v", err) + } + hasExecDir := false + for _, arg := range args { + if arg == execDir { + hasExecDir = true + break + } + } + if !hasExecDir { + t.Fatalf("buildLinuxBwrapArgs() missing resolved chdir: %v", args) + } + for i := range args { + if args[i] == "--" { + if i+1 >= len(args) || args[i+1] != resolvedPath { + t.Fatalf("buildLinuxBwrapArgs() exec path = %v, want %q after --", args, resolvedPath) + } + return + } + } + t.Fatalf("buildLinuxBwrapArgs() missing exec delimiter: %v", args) +} + +func TestAppendLinuxArgumentMounts_AddsAbsoluteArgumentPaths(t *testing.T) { + root := t.TempDir() + input := filepath.Join(root, "input.txt") + if err := os.WriteFile(input, []byte("data"), 0o644); err != nil { + t.Fatal(err) + } + output := filepath.Join(root, "out", "result.txt") + if err := os.MkdirAll(filepath.Dir(output), 0o755); err != nil { + t.Fatal(err) + } + + plan := appendLinuxArgumentMounts(nil, []string{input, "--output=" + output}) + if len(plan) != 2 { + t.Fatalf("appendLinuxArgumentMounts() len = %d, want 2", len(plan)) + } + if plan[0].Source != input || plan[0].Mode != "ro" { + t.Fatalf("appendLinuxArgumentMounts()[0] = %+v, want source=%q mode=ro", plan[0], input) + } + if plan[1].Source != filepath.Dir(output) || plan[1].Mode != "rw" { + t.Fatalf("appendLinuxArgumentMounts()[1] = %+v, want source=%q mode=rw", plan[1], filepath.Dir(output)) + } +} diff --git a/pkg/isolation/platform_other.go b/pkg/isolation/platform_other.go new file mode 100644 index 000000000..d8d06e2ec --- /dev/null +++ b/pkg/isolation/platform_other.go @@ -0,0 +1,22 @@ +//go:build !linux && !windows + +package isolation + +import ( + "os/exec" + + "github.com/sipeed/picoclaw/pkg/config" +) + +func applyPlatformIsolation(cmd *exec.Cmd, isolation config.IsolationConfig, root string) error { + // Unsupported platforms currently keep the command unchanged. Callers rely on + // Preflight and higher-level checks to surface unsupported isolation modes. + return nil +} + +func postStartPlatformIsolation(cmd *exec.Cmd, isolation config.IsolationConfig, root string) error { + return nil +} + +func cleanupPendingPlatformResources(cmd *exec.Cmd) { +} diff --git a/pkg/isolation/platform_windows.go b/pkg/isolation/platform_windows.go new file mode 100644 index 000000000..9434976f7 --- /dev/null +++ b/pkg/isolation/platform_windows.go @@ -0,0 +1,217 @@ +//go:build windows + +package isolation + +import ( + "fmt" + "os/exec" + "sync" + "syscall" + "unsafe" + + "golang.org/x/sys/windows" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" +) + +const disableMaxPrivilege = 0x1 + +// windowsProcessResources holds native handles that must live for the lifetime +// of an isolated child process. +type windowsProcessResources struct { + job windows.Handle + token windows.Token +} + +var ( + windowsProcessResourcesByPID sync.Map + windowsPendingResources sync.Map + advapi32 = windows.NewLazySystemDLL("advapi32.dll") + procCreateRestrictedToken = advapi32.NewProc("CreateRestrictedToken") +) + +func applyPlatformIsolation(cmd *exec.Cmd, isolation config.IsolationConfig, root string) error { + if !isolation.Enabled || cmd == nil { + return nil + } + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + rules := BuildWindowsAccessRules(root, isolation.ExposePaths) + logger.InfoCF("isolation", "windows isolation process constraints", + map[string]any{ + "root": root, + "command": cmd.Path, + "rules": formatWindowsAccessRules(rules), + "note": "Windows currently enforces restricted token, low integrity, and job object limits; expose_paths filesystem remapping is rejected during preflight", + }) + // Create the restricted token before the process starts so CreateProcess uses + // the reduced privilege set from the first instruction. + restrictedToken, err := createRestrictedPrimaryToken() + if err != nil { + return fmt.Errorf("create restricted primary token: %w", err) + } + cmd.SysProcAttr.CreationFlags |= windows.CREATE_NEW_PROCESS_GROUP | windows.CREATE_BREAKAWAY_FROM_JOB + cmd.SysProcAttr.Token = syscall.Token(restrictedToken) + windowsPendingResources.Store(cmd, windowsProcessResources{token: restrictedToken}) + return nil +} + +func postStartPlatformIsolation(cmd *exec.Cmd, isolation config.IsolationConfig, root string) error { + if !isolation.Enabled || cmd == nil || cmd.Process == nil { + return nil + } + resourcesAny, _ := windowsPendingResources.LoadAndDelete(cmd) + resources, _ := resourcesAny.(windowsProcessResources) + // Job objects can only be attached after the process exists, so the Windows + // backend finishes isolation in this post-start hook. + job, err := windows.CreateJobObject(nil, nil) + if err != nil { + if resources.token != 0 { + _ = resources.token.Close() + } + return fmt.Errorf("create windows job object: %w", err) + } + + info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{} + info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + if _, err := windows.SetInformationJobObject( + job, + windows.JobObjectExtendedLimitInformation, + uintptr(unsafe.Pointer(&info)), + uint32(unsafe.Sizeof(info)), + ); err != nil { + _ = windows.CloseHandle(job) + if resources.token != 0 { + _ = resources.token.Close() + } + return fmt.Errorf("set windows job object info: %w", err) + } + + proc, err := windows.OpenProcess( + windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE|windows.PROCESS_QUERY_LIMITED_INFORMATION|windows.SYNCHRONIZE, + false, + uint32(cmd.Process.Pid), + ) + if err != nil { + _ = windows.CloseHandle(job) + if resources.token != 0 { + _ = resources.token.Close() + } + return fmt.Errorf("open process for job assignment: %w", err) + } + + if err := windows.AssignProcessToJobObject(job, proc); err != nil { + _ = windows.CloseHandle(proc) + _ = windows.CloseHandle(job) + if resources.token != 0 { + _ = resources.token.Close() + } + return fmt.Errorf("assign process to job object: %w", err) + } + + if resources.token != 0 { + _ = resources.token.Close() + } + resources.job = job + windowsProcessResourcesByPID.Store(cmd.Process.Pid, resources) + go reapWindowsProcessResources(cmd.Process.Pid, proc, job) + return nil +} + +func cleanupPendingPlatformResources(cmd *exec.Cmd) { + if cmd == nil { + return + } + resourcesAny, ok := windowsPendingResources.LoadAndDelete(cmd) + if !ok { + return + } + resources, _ := resourcesAny.(windowsProcessResources) + if resources.token != 0 { + _ = resources.token.Close() + } +} + +func reapWindowsProcessResources(pid int, proc windows.Handle, job windows.Handle) { + _, _ = windows.WaitForSingleObject(proc, windows.INFINITE) + _ = windows.CloseHandle(proc) + _ = windows.CloseHandle(job) + windowsProcessResourcesByPID.Delete(pid) +} + +// createRestrictedPrimaryToken duplicates the current process token, removes +// maximum privileges, and lowers integrity before it is assigned to a child. +func createRestrictedPrimaryToken() (windows.Token, error) { + var current windows.Token + if err := windows.OpenProcessToken( + windows.CurrentProcess(), + windows.TOKEN_DUPLICATE|windows.TOKEN_ASSIGN_PRIMARY|windows.TOKEN_QUERY|windows.TOKEN_ADJUST_DEFAULT, + ¤t, + ); err != nil { + return 0, err + } + defer current.Close() + + var restricted windows.Token + r1, _, e1 := procCreateRestrictedToken.Call( + uintptr(current), + uintptr(disableMaxPrivilege), + 0, + 0, + 0, + 0, + 0, + 0, + 0, + uintptr(unsafe.Pointer(&restricted)), + ) + if r1 == 0 { + if e1 != nil && e1 != syscall.Errno(0) { + return 0, e1 + } + return 0, syscall.EINVAL + } + if err := setTokenLowIntegrity(restricted); err != nil { + _ = restricted.Close() + return 0, err + } + return restricted, nil +} + +// setTokenLowIntegrity lowers the token integrity level so writes to higher +// integrity locations are blocked by the OS. +func setTokenLowIntegrity(token windows.Token) error { + lowSID, err := windows.CreateWellKnownSid(windows.WinLowLabelSid) + if err != nil { + return fmt.Errorf("create low integrity sid: %w", err) + } + tml := windows.Tokenmandatorylabel{ + Label: windows.SIDAndAttributes{ + Sid: lowSID, + Attributes: windows.SE_GROUP_INTEGRITY, + }, + } + if err := windows.SetTokenInformation( + token, + windows.TokenIntegrityLevel, + (*byte)(unsafe.Pointer(&tml)), + tml.Size(), + ); err != nil { + return fmt.Errorf("set token low integrity: %w", err) + } + return nil +} + +// formatWindowsAccessRules reshapes the internal rules for structured logging. +func formatWindowsAccessRules(rules []AccessRule) []map[string]string { + formatted := make([]map[string]string, 0, len(rules)) + for _, rule := range rules { + formatted = append(formatted, map[string]string{ + "path": rule.Path, + "mode": rule.Mode, + }) + } + return formatted +} diff --git a/pkg/isolation/runtime.go b/pkg/isolation/runtime.go new file mode 100644 index 000000000..b2de98b88 --- /dev/null +++ b/pkg/isolation/runtime.go @@ -0,0 +1,443 @@ +package isolation + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "sync" + + "github.com/sipeed/picoclaw/pkg" + "github.com/sipeed/picoclaw/pkg/config" +) + +// MountRule describes a source-to-target mount exposed inside the Linux +// isolation view. +type MountRule struct { + Source string + Target string + Mode string +} + +// AccessRule describes the effective Windows-side access rule for a host path. +type AccessRule struct { + Path string + Mode string +} + +// UserEnv contains the redirected per-instance user directories injected into +// isolated child processes. +type UserEnv struct { + Home string + Tmp string + Config string + Cache string + State string + AppData string + LocalAppData string +} + +var ( + isolationMu sync.RWMutex + currentIsolation = config.DefaultConfig().Isolation +) + +// Configure updates the process-wide isolation state used by subsequent child +// process launches. +func Configure(cfg *config.Config) { + isolationMu.Lock() + defer isolationMu.Unlock() + if cfg == nil { + defaults := config.DefaultConfig() + currentIsolation = defaults.Isolation + return + } + currentIsolation = cfg.Isolation +} + +// CurrentConfig returns the currently active isolation settings. +func CurrentConfig() config.IsolationConfig { + isolationMu.RLock() + defer isolationMu.RUnlock() + return currentIsolation +} + +// ResolveInstanceRoot resolves the instance root used to build the isolated +// filesystem and redirected user environment. +func ResolveInstanceRoot() (string, error) { + root := filepath.Clean(config.GetHome()) + if root == "." { + return "", fmt.Errorf("instance root resolved to current directory") + } + return root, nil +} + +// PrepareInstanceRoot creates the directories required by the isolation runtime. +func PrepareInstanceRoot(root string) error { + for _, dir := range InstanceDirs(root) { + if err := os.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("prepare instance dir %s: %w", dir, err) + } + } + return nil +} + +// InstanceDirs returns the directories that must exist under the instance root +// for isolation-aware child processes. +func InstanceDirs(root string) []string { + dirs := []string{ + root, + filepath.Join(root, "skills"), + filepath.Join(root, "logs"), + filepath.Join(root, "cache"), + filepath.Join(root, "state"), + filepath.Join(root, "runtime-user-env"), + filepath.Join(root, "runtime-user-env", "home"), + filepath.Join(root, "runtime-user-env", "tmp"), + filepath.Join(root, "runtime-user-env", "config"), + filepath.Join(root, "runtime-user-env", "cache"), + filepath.Join(root, "runtime-user-env", "state"), + } + dirs = append(dirs, filepath.Join(root, pkg.WorkspaceName)) + if runtime.GOOS == "windows" { + dirs = append(dirs, + filepath.Join(root, "runtime-user-env", "AppData", "Roaming"), + filepath.Join(root, "runtime-user-env", "AppData", "Local"), + ) + } + return dirs +} + +// ResolveUserEnv derives the redirected user directories rooted under the +// instance runtime area. +func ResolveUserEnv(root string) UserEnv { + base := filepath.Join(root, "runtime-user-env") + return UserEnv{ + Home: filepath.Join(base, "home"), + Tmp: filepath.Join(base, "tmp"), + Config: filepath.Join(base, "config"), + Cache: filepath.Join(base, "cache"), + State: filepath.Join(base, "state"), + AppData: filepath.Join(base, "AppData", "Roaming"), + LocalAppData: filepath.Join(base, "AppData", "Local"), + } +} + +// ApplyUserEnv rewrites the child process environment so home, temp, and +// platform-specific user-data directories point into the instance root. +func ApplyUserEnv(cmd *exec.Cmd, root string) { + userEnv := ResolveUserEnv(root) + envMap := make(map[string]string) + for _, item := range cmd.Environ() { + if idx := strings.IndexRune(item, '='); idx > 0 { + envMap[item[:idx]] = item[idx+1:] + } + } + + if runtime.GOOS == "windows" { + envMap["USERPROFILE"] = userEnv.Home + envMap["HOME"] = userEnv.Home + envMap["TEMP"] = userEnv.Tmp + envMap["TMP"] = userEnv.Tmp + envMap["APPDATA"] = userEnv.AppData + envMap["LOCALAPPDATA"] = userEnv.LocalAppData + } else { + envMap["HOME"] = userEnv.Home + envMap["TMPDIR"] = userEnv.Tmp + envMap["XDG_CONFIG_HOME"] = userEnv.Config + envMap["XDG_CACHE_HOME"] = userEnv.Cache + envMap["XDG_STATE_HOME"] = userEnv.State + } + + env := make([]string, 0, len(envMap)) + for k, v := range envMap { + env = append(env, fmt.Sprintf("%s=%s", k, v)) + } + cmd.Env = env +} + +// ValidateExposePaths verifies the user-supplied path exposure rules before a +// child process is started. +func ValidateExposePaths(items []config.ExposePath) error { + seen := map[string]struct{}{} + for _, item := range items { + if item.Source == "" { + return fmt.Errorf("source is required") + } + if item.Mode != "ro" && item.Mode != "rw" { + return fmt.Errorf("invalid expose_paths mode: %s", item.Mode) + } + + source := filepath.Clean(item.Source) + target := item.Target + if target == "" { + target = source + } + target = filepath.Clean(target) + + if !filepath.IsAbs(source) || !filepath.IsAbs(target) { + return fmt.Errorf("source and target must be absolute paths") + } + if _, ok := seen[target]; ok { + return fmt.Errorf("duplicate expose_path target: %s", target) + } + seen[target] = struct{}{} + } + return nil +} + +// NormalizeExposePath fills implicit defaults and cleans path values so merge +// and validation logic can work with canonical paths. +func NormalizeExposePath(item config.ExposePath) config.ExposePath { + source := filepath.Clean(item.Source) + target := item.Target + if target == "" { + target = source + } + return config.ExposePath{ + Source: source, + Target: filepath.Clean(target), + Mode: item.Mode, + } +} + +// DefaultExposePaths returns the minimum built-in host paths required for the +// current platform to run isolated child processes. +func DefaultExposePaths(root string) []config.ExposePath { + items := []config.ExposePath{{ + Source: root, + Target: root, + Mode: "rw", + }} + if runtime.GOOS == "linux" { + items = append(items, defaultLinuxSystemExposePaths()...) + } + return items +} + +func defaultLinuxSystemExposePaths() []config.ExposePath { + return existingExposePaths([]config.ExposePath{ + {Source: "/usr", Target: "/usr", Mode: "ro"}, + {Source: "/bin", Target: "/bin", Mode: "ro"}, + {Source: "/lib", Target: "/lib", Mode: "ro"}, + {Source: "/lib64", Target: "/lib64", Mode: "ro"}, + {Source: "/etc/resolv.conf", Target: "/etc/resolv.conf", Mode: "ro"}, + {Source: "/etc/hosts", Target: "/etc/hosts", Mode: "ro"}, + {Source: "/etc/nsswitch.conf", Target: "/etc/nsswitch.conf", Mode: "ro"}, + {Source: "/etc/passwd", Target: "/etc/passwd", Mode: "ro"}, + {Source: "/etc/group", Target: "/etc/group", Mode: "ro"}, + {Source: "/etc/ssl", Target: "/etc/ssl", Mode: "ro"}, + {Source: "/etc/pki", Target: "/etc/pki", Mode: "ro"}, + {Source: "/etc/ca-certificates", Target: "/etc/ca-certificates", Mode: "ro"}, + {Source: "/usr/share/ca-certificates", Target: "/usr/share/ca-certificates", Mode: "ro"}, + {Source: "/usr/local/share/ca-certificates", Target: "/usr/local/share/ca-certificates", Mode: "ro"}, + {Source: "/etc/alternatives", Target: "/etc/alternatives", Mode: "ro"}, + {Source: "/usr/share/zoneinfo", Target: "/usr/share/zoneinfo", Mode: "ro"}, + {Source: "/etc/localtime", Target: "/etc/localtime", Mode: "ro"}, + }) +} + +// existingExposePaths keeps only the builtin host paths that exist on the +// current machine so Linux isolation does not fail on distro-specific paths. +func existingExposePaths(items []config.ExposePath) []config.ExposePath { + filtered := make([]config.ExposePath, 0, len(items)) + for _, item := range items { + if _, err := os.Stat(item.Source); err == nil { + filtered = append(filtered, item) + } + } + return filtered +} + +// MergeExposePaths merges built-in rules with user overrides. Rules are keyed +// by target path so later entries replace earlier ones for the same target. +func MergeExposePaths(defaults []config.ExposePath, overrides []config.ExposePath) []config.ExposePath { + merged := make([]config.ExposePath, 0, len(defaults)+len(overrides)) + indexByTarget := make(map[string]int, len(defaults)+len(overrides)) + appendOrReplace := func(item config.ExposePath) { + normalized := NormalizeExposePath(item) + if idx, ok := indexByTarget[normalized.Target]; ok { + merged[idx] = normalized + return + } + indexByTarget[normalized.Target] = len(merged) + merged = append(merged, normalized) + } + for _, item := range defaults { + appendOrReplace(item) + } + for _, item := range overrides { + appendOrReplace(item) + } + return merged +} + +// BuildLinuxMountPlan converts the merged expose-path configuration into the +// mount rules consumed by the Linux bubblewrap backend. +func BuildLinuxMountPlan(root string, overrides []config.ExposePath) []MountRule { + merged := MergeExposePaths(DefaultExposePaths(root), overrides) + plan := make([]MountRule, 0, len(merged)) + for _, item := range merged { + plan = append(plan, MountRule{Source: item.Source, Target: item.Target, Mode: item.Mode}) + } + return plan +} + +// BuildWindowsAccessRules derives the host-path access policy used by the +// Windows restricted-token backend. +func BuildWindowsAccessRules(root string, overrides []config.ExposePath) []AccessRule { + merged := MergeExposePaths(nil, overrides) + rules := make([]AccessRule, 0, len(merged)+1) + rules = append(rules, AccessRule{Path: root, Mode: "rw"}) + for _, item := range merged { + rules = append(rules, AccessRule{Path: item.Source, Mode: item.Mode}) + } + return rules +} + +func validateWindowsExposePaths(items []config.ExposePath) error { + if len(items) == 0 { + return nil + } + return fmt.Errorf("windows isolation does not yet support expose_paths filesystem rules") +} + +// IsSupported reports whether the current platform has an implemented isolation +// backend. +func IsSupported() bool { + return isSupportedOn(runtime.GOOS) +} + +func isSupportedOn(goos string) bool { + switch goos { + case "linux", "windows": + return true + default: + return false + } +} + +// Preflight validates the configured isolation state and prepares the instance +// runtime directories before any child process is launched. +func Preflight() error { + isolation := CurrentConfig() + if !isolation.Enabled { + return nil + } + if !IsSupported() { + return fmt.Errorf("subprocess isolation is not supported on %s", runtime.GOOS) + } + root, err := ResolveInstanceRoot() + if err != nil { + return err + } + if err := PrepareInstanceRoot(root); err != nil { + return err + } + if err := ValidateExposePaths(isolation.ExposePaths); err != nil { + return err + } + if runtime.GOOS == "linux" { + for _, rule := range BuildLinuxMountPlan(root, isolation.ExposePaths) { + if rule.Source == "" || rule.Target == "" { + return fmt.Errorf("invalid linux mount rule") + } + } + } + if runtime.GOOS == "windows" { + if err := validateWindowsExposePaths(isolation.ExposePaths); err != nil { + return err + } + for _, rule := range BuildWindowsAccessRules(root, isolation.ExposePaths) { + if rule.Path == "" { + return fmt.Errorf("invalid windows access rule") + } + } + } + return nil +} + +// Start prepares isolation for the command, starts it, and applies any +// post-start platform hooks required by the active backend. +func Start(cmd *exec.Cmd) error { + if err := PrepareCommand(cmd); err != nil { + return err + } + if err := cmd.Start(); err != nil { + cleanupPendingPlatformResources(cmd) + return err + } + isolation := CurrentConfig() + root := "" + if isolation.Enabled { + var err error + root, err = ResolveInstanceRoot() + if err != nil { + terminateStartedCommand(cmd) + return err + } + } + if err := postStartPlatformIsolation(cmd, isolation, root); err != nil { + terminateStartedCommand(cmd) + return err + } + return nil +} + +// Run is the Start-and-Wait helper that keeps the same isolation behavior as +// Start while returning the command's final exit status. +func Run(cmd *exec.Cmd) error { + if err := PrepareCommand(cmd); err != nil { + return err + } + if err := cmd.Start(); err != nil { + cleanupPendingPlatformResources(cmd) + return err + } + isolation := CurrentConfig() + root := "" + if isolation.Enabled { + var err error + root, err = ResolveInstanceRoot() + if err != nil { + terminateStartedCommand(cmd) + return err + } + } + if err := postStartPlatformIsolation(cmd, isolation, root); err != nil { + terminateStartedCommand(cmd) + return err + } + return cmd.Wait() +} + +func terminateStartedCommand(cmd *exec.Cmd) { + cleanupPendingPlatformResources(cmd) + if cmd == nil || cmd.Process == nil { + return + } + _ = cmd.Process.Kill() + _ = cmd.Wait() +} + +// PrepareCommand mutates the command in-place so it inherits the configured +// isolated environment before being started by the caller. +func PrepareCommand(cmd *exec.Cmd) error { + isolation := CurrentConfig() + if err := Preflight(); err != nil { + return err + } + if isolation.Enabled { + root, err := ResolveInstanceRoot() + if err != nil { + return err + } + ApplyUserEnv(cmd, root) + if err := applyPlatformIsolation(cmd, isolation, root); err != nil { + return err + } + } + return nil +} diff --git a/pkg/isolation/runtime_test.go b/pkg/isolation/runtime_test.go new file mode 100644 index 000000000..213c4b065 --- /dev/null +++ b/pkg/isolation/runtime_test.go @@ -0,0 +1,245 @@ +package isolation + +import ( + "os" + "os/exec" + "path/filepath" + "runtime" + "testing" + + "github.com/sipeed/picoclaw/pkg" + "github.com/sipeed/picoclaw/pkg/config" +) + +func TestResolveInstanceRoot_UsesPicoclawHome(t *testing.T) { + t.Setenv(config.EnvHome, "/custom/picoclaw/home") + root, err := ResolveInstanceRoot() + if err != nil { + t.Fatalf("ResolveInstanceRoot() error = %v", err) + } + if root != "/custom/picoclaw/home" { + t.Fatalf("ResolveInstanceRoot() = %q, want %q", root, "/custom/picoclaw/home") + } +} + +func TestPrepareInstanceRoot_CreatesDirectories(t *testing.T) { + root := filepath.Join(t.TempDir(), "instance") + if err := PrepareInstanceRoot(root); err != nil { + t.Fatalf("PrepareInstanceRoot() error = %v", err) + } + for _, dir := range InstanceDirs(root) { + if info, err := os.Stat(dir); err != nil { + t.Fatalf("os.Stat(%q): %v", dir, err) + } else if !info.IsDir() { + t.Fatalf("%q is not a directory", dir) + } + } +} + +func TestInstanceDirs_UsesInstanceWorkspaceNotGlobalState(t *testing.T) { + root := filepath.Join(t.TempDir(), "instance") + cfg := config.DefaultConfig() + cfg.Isolation.Enabled = true + cfg.Agents.Defaults.Workspace = filepath.Join(t.TempDir(), "external-workspace") + Configure(cfg) + t.Cleanup(func() { Configure(config.DefaultConfig()) }) + + dirs := InstanceDirs(root) + wantWorkspace := filepath.Join(root, pkg.WorkspaceName) + found := false + for _, dir := range dirs { + if dir == wantWorkspace { + found = true + } + if dir == cfg.WorkspacePath() { + t.Fatalf("InstanceDirs() should not depend on process-wide workspace state: %q", dir) + } + } + if !found { + t.Fatalf("InstanceDirs() missing instance workspace dir %q", wantWorkspace) + } +} + +func TestIsSupportedOn(t *testing.T) { + tests := []struct { + goos string + want bool + }{ + {goos: "linux", want: true}, + {goos: "windows", want: true}, + {goos: "darwin", want: false}, + {goos: "freebsd", want: false}, + } + for _, tt := range tests { + if got := isSupportedOn(tt.goos); got != tt.want { + t.Fatalf("isSupportedOn(%q) = %v, want %v", tt.goos, got, tt.want) + } + } +} + +func TestValidateExposePaths(t *testing.T) { + err := ValidateExposePaths([]config.ExposePath{{Source: "/src", Target: "/dst", Mode: "ro"}}) + if err != nil { + t.Fatalf("ValidateExposePaths() error = %v", err) + } + + err = ValidateExposePaths([]config.ExposePath{{Source: "/src", Target: "/dst", Mode: "bad"}}) + if err == nil { + t.Fatal("ValidateExposePaths() expected invalid mode error") + } + + err = ValidateExposePaths( + []config.ExposePath{ + {Source: "/src", Target: "/dst", Mode: "ro"}, + {Source: "/other", Target: "/dst", Mode: "rw"}, + }, + ) + if err == nil { + t.Fatal("ValidateExposePaths() expected duplicate target error") + } +} + +func TestMergeExposePaths_OverrideByTarget(t *testing.T) { + merged := MergeExposePaths( + []config.ExposePath{{Source: "/src-a", Target: "/dst", Mode: "ro"}}, + []config.ExposePath{{Source: "/src-b", Target: "/dst", Mode: "rw"}}, + ) + if len(merged) != 1 { + t.Fatalf("MergeExposePaths len = %d, want 1", len(merged)) + } + if got := merged[0]; got.Source != "/src-b" || got.Target != "/dst" || got.Mode != "rw" { + t.Fatalf("merged[0] = %+v, want source=/src-b target=/dst mode=rw", got) + } +} + +func TestBuildLinuxMountPlan(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("linux-only default mount set") + } + plan := BuildLinuxMountPlan("/rootdir", []config.ExposePath{{Source: "/src", Target: "/dst", Mode: "ro"}}) + if len(plan) == 0 { + t.Fatal("BuildLinuxMountPlan returned empty plan") + } + foundRoot := false + foundOverride := false + for _, rule := range plan { + if rule.Source == "/rootdir" && rule.Target == "/rootdir" && rule.Mode == "rw" { + foundRoot = true + } + if rule.Source == "/src" && rule.Target == "/dst" && rule.Mode == "ro" { + foundOverride = true + } + } + if !foundRoot { + t.Fatal("BuildLinuxMountPlan missing root mapping") + } + if !foundOverride { + t.Fatal("BuildLinuxMountPlan missing override mapping") + } +} + +func TestBuildWindowsAccessRules(t *testing.T) { + rules := BuildWindowsAccessRules( + `C:\picoclaw`, + []config.ExposePath{{Source: `D:\data`, Target: `C:\mapped`, Mode: "ro"}}, + ) + if len(rules) == 0 { + t.Fatal("BuildWindowsAccessRules returned empty rules") + } + foundRoot := false + foundOverride := false + for _, rule := range rules { + if rule.Path == `C:\picoclaw` && rule.Mode == "rw" { + foundRoot = true + } + if rule.Path == `D:\data` && rule.Mode == "ro" { + foundOverride = true + } + } + if !foundRoot { + t.Fatal("BuildWindowsAccessRules missing root rule") + } + if !foundOverride { + t.Fatal("BuildWindowsAccessRules missing override rule") + } +} + +func TestValidateWindowsExposePaths(t *testing.T) { + if err := validateWindowsExposePaths(nil); err != nil { + t.Fatalf("validateWindowsExposePaths(nil) error = %v", err) + } + err := validateWindowsExposePaths([]config.ExposePath{{Source: `D:\data`, Target: `D:\data`, Mode: "ro"}}) + if err == nil { + t.Fatal("validateWindowsExposePaths() expected error for expose_paths") + } +} + +func TestDefaultLinuxSystemExposePaths(t *testing.T) { + paths := defaultLinuxSystemExposePaths() + needed := map[string]bool{} + for _, path := range []string{"/etc/hosts", "/etc/nsswitch.conf", "/etc/ssl", "/usr/share/zoneinfo", "/etc/localtime"} { + if _, err := os.Stat(path); err == nil { + needed[path] = false + } + } + for _, item := range paths { + if _, ok := needed[item.Source]; ok { + needed[item.Source] = true + } + } + for path, found := range needed { + if !found { + t.Fatalf("defaultLinuxSystemExposePaths missing %s", path) + } + } +} + +func TestExistingExposePaths_SkipsMissingPaths(t *testing.T) { + existing := filepath.Join(t.TempDir(), "existing") + if err := os.MkdirAll(existing, 0o755); err != nil { + t.Fatalf("os.MkdirAll() error = %v", err) + } + filtered := existingExposePaths([]config.ExposePath{ + {Source: existing, Target: existing, Mode: "ro"}, + {Source: filepath.Join(t.TempDir(), "missing"), Target: "/missing", Mode: "ro"}, + }) + if len(filtered) != 1 { + t.Fatalf("existingExposePaths() len = %d, want 1", len(filtered)) + } + if got := filtered[0]; got.Source != existing { + t.Fatalf("existingExposePaths()[0] = %+v, want source=%q", got, existing) + } +} + +func TestPrepareCommand_AppliesUserEnv(t *testing.T) { + t.Setenv(config.EnvHome, filepath.Join(t.TempDir(), "home")) + if runtime.GOOS == "linux" { + binDir := filepath.Join(t.TempDir(), "bin") + if err := os.MkdirAll(binDir, 0o755); err != nil { + t.Fatalf("os.MkdirAll() error = %v", err) + } + fakeBwrap := filepath.Join(binDir, "bwrap") + if err := os.WriteFile(fakeBwrap, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil { + t.Fatalf("os.WriteFile() error = %v", err) + } + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + } + cfg := config.DefaultConfig() + cfg.Isolation.Enabled = true + Configure(cfg) + t.Cleanup(func() { Configure(config.DefaultConfig()) }) + cmd := exec.Command("sh", "-c", "true") + if err := PrepareCommand(cmd); err != nil { + t.Fatalf("PrepareCommand() error = %v", err) + } + hasHome := false + for _, env := range cmd.Env { + if len(env) > 5 && env[:5] == "HOME=" { + hasHome = true + break + } + } + if runtime.GOOS != "windows" && !hasHome { + t.Fatal("PrepareCommand() did not inject HOME") + } +} diff --git a/pkg/mcp/isolated_command_transport.go b/pkg/mcp/isolated_command_transport.go new file mode 100644 index 000000000..f54b4af8b --- /dev/null +++ b/pkg/mcp/isolated_command_transport.go @@ -0,0 +1,226 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os/exec" + "sync" + "syscall" + "time" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/sipeed/picoclaw/pkg/isolation" +) + +var isolatedCommandTerminateDuration = 5 * time.Second + +// isolatedCommandTransport mirrors the SDK command transport but routes +// process startup through pkg/isolation so Windows post-start hooks run too. +type isolatedCommandTransport struct { + Command *exec.Cmd + TerminateDuration time.Duration +} + +func (t *isolatedCommandTransport) Connect(ctx context.Context) (sdkmcp.Connection, error) { + stdout, err := t.Command.StdoutPipe() + if err != nil { + return nil, err + } + stdout = io.NopCloser(stdout) + stdin, err := t.Command.StdinPipe() + if err != nil { + return nil, err + } + if err := isolation.Start(t.Command); err != nil { + return nil, err + } + td := t.TerminateDuration + if td <= 0 { + td = isolatedCommandTerminateDuration + } + return newIsolatedIOConn(&isolatedPipeRWC{cmd: t.Command, stdout: stdout, stdin: stdin, terminateDuration: td}), nil +} + +type isolatedPipeRWC struct { + cmd *exec.Cmd + stdout io.ReadCloser + stdin io.WriteCloser + terminateDuration time.Duration +} + +func (s *isolatedPipeRWC) Read(p []byte) (n int, err error) { + return s.stdout.Read(p) +} + +func (s *isolatedPipeRWC) Write(p []byte) (n int, err error) { + return s.stdin.Write(p) +} + +func (s *isolatedPipeRWC) Close() error { + if err := s.stdin.Close(); err != nil { + return fmt.Errorf("closing stdin: %v", err) + } + resChan := make(chan error, 1) + go func() { + resChan <- s.cmd.Wait() + }() + wait := func() (error, bool) { + select { + case err := <-resChan: + return err, true + case <-time.After(s.terminateDuration): + } + return nil, false + } + if err, ok := wait(); ok { + return err + } + if err := s.cmd.Process.Signal(syscall.SIGTERM); err == nil { + if err, ok := wait(); ok { + return err + } + } + if err := s.cmd.Process.Kill(); err != nil { + return err + } + if err, ok := wait(); ok { + return err + } + return fmt.Errorf("unresponsive subprocess") +} + +type isolatedIOConn struct { + writeMu sync.Mutex + rwc io.ReadWriteCloser + incoming <-chan isolatedMsgOrErr + queue []jsonrpc.Message + closeOnce sync.Once + closed chan struct{} + closeErr error +} + +type isolatedMsgOrErr struct { + msg json.RawMessage + err error +} + +func newIsolatedIOConn(rwc io.ReadWriteCloser) *isolatedIOConn { + incoming := make(chan isolatedMsgOrErr) + closed := make(chan struct{}) + go func() { + dec := json.NewDecoder(rwc) + for { + var raw json.RawMessage + err := dec.Decode(&raw) + if err == nil { + var tr [1]byte + if n, readErr := dec.Buffered().Read(tr[:]); n > 0 { + if tr[0] != '\n' && tr[0] != '\r' { + err = fmt.Errorf("invalid trailing data at the end of stream") + } + } else if readErr != nil && readErr != io.EOF { + err = readErr + } + } + select { + case incoming <- isolatedMsgOrErr{msg: raw, err: err}: + case <-closed: + return + } + if err != nil { + return + } + } + }() + return &isolatedIOConn{rwc: rwc, incoming: incoming, closed: closed} +} + +func (c *isolatedIOConn) SessionID() string { return "" } + +func (c *isolatedIOConn) Read(ctx context.Context) (jsonrpc.Message, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + if len(c.queue) > 0 { + next := c.queue[0] + c.queue = c.queue[1:] + return next, nil + } + var raw json.RawMessage + select { + case <-ctx.Done(): + return nil, ctx.Err() + case v := <-c.incoming: + if v.err != nil { + return nil, v.err + } + raw = v.msg + case <-c.closed: + return nil, io.EOF + } + msgs, err := readIsolatedBatch(raw) + if err != nil { + return nil, err + } + c.queue = msgs[1:] + return msgs[0], nil +} + +func readIsolatedBatch(data []byte) ([]jsonrpc.Message, error) { + var rawBatch []json.RawMessage + if err := json.Unmarshal(data, &rawBatch); err == nil { + if len(rawBatch) == 0 { + return nil, fmt.Errorf("empty batch") + } + msgs := make([]jsonrpc.Message, 0, len(rawBatch)) + for _, raw := range rawBatch { + msg, err := jsonrpc.DecodeMessage(raw) + if err != nil { + return nil, err + } + msgs = append(msgs, msg) + } + return msgs, nil + } + msg, err := jsonrpc.DecodeMessage(data) + if err != nil { + return nil, err + } + return []jsonrpc.Message{msg}, nil +} + +func (c *isolatedIOConn) Write(ctx context.Context, msg jsonrpc.Message) error { + select { + case <-ctx.Done(): + return ctx.Err() + default: + } + c.writeMu.Lock() + defer c.writeMu.Unlock() + data, err := jsonrpc.EncodeMessage(msg) + if err != nil { + return fmt.Errorf("marshaling message: %v", err) + } + data = append(data, '\n') + _, err = c.rwc.Write(data) + return err +} + +func (c *isolatedIOConn) Close() error { + c.closeOnce.Do(func() { + c.closeErr = c.rwc.Close() + close(c.closed) + }) + return c.closeErr +} + +var ( + _ sdkmcp.Transport = (*isolatedCommandTransport)(nil) + _ sdkmcp.Connection = (*isolatedIOConn)(nil) +) diff --git a/pkg/mcp/manager.go b/pkg/mcp/manager.go index 323df0312..f589f82a9 100644 --- a/pkg/mcp/manager.go +++ b/pkg/mcp/manager.go @@ -365,8 +365,7 @@ func (m *Manager) ConnectServer( env = append(env, fmt.Sprintf("%s=%s", k, v)) } cmd.Env = env - - transport = &mcp.CommandTransport{Command: cmd} + transport = &isolatedCommandTransport{Command: cmd} default: return fmt.Errorf( "unsupported transport type: %s (supported: stdio, sse, http)", diff --git a/pkg/pid/pidfile.go b/pkg/pid/pidfile.go index 69d02bc65..f7c1f42b2 100644 --- a/pkg/pid/pidfile.go +++ b/pkg/pid/pidfile.go @@ -4,6 +4,7 @@ import ( "crypto/rand" "encoding/hex" "encoding/json" + "errors" "fmt" "os" "path/filepath" @@ -16,6 +17,8 @@ import ( const pidFileName = ".picoclaw.pid" +var errInvalidPidFile = errors.New("invalid pid file") + // PidFileData is the JSON structure stored in the PID file. type PidFileData struct { PID int `json:"pid"` @@ -109,6 +112,14 @@ func ReadPidFileWithCheck(homePath string) *PidFileData { pidPath := pidFilePath(homePath) data, err := readPidFileUnlocked(pidPath) if err != nil { + if os.IsNotExist(err) { + return nil + } + if errors.Is(err, errInvalidPidFile) { + logger.Warnf("invalid pid file, remove it: %s (%v)", pidPath, err) + _ = os.Remove(pidPath) + return nil + } logger.Debugf("failed to read pid file: %s", err) return nil } @@ -140,6 +151,30 @@ func RemovePidFile(homePath string) { os.Remove(pidPath) } +// RemovePidFileIfPID deletes the PID file only when the recorded PID matches +// expectedPID. It returns true when the file is removed successfully. +func RemovePidFileIfPID(homePath string, expectedPID int) bool { + if expectedPID <= 0 { + return false + } + + pidMu.Lock() + defer pidMu.Unlock() + + pidPath := pidFilePath(homePath) + data, err := readPidFileUnlocked(pidPath) + if err != nil { + return false + } + if data.PID != expectedPID { + return false + } + if err := os.Remove(pidPath); err != nil { + return false + } + return true +} + // readPidFileUnlocked reads the PID file without acquiring the lock. // Caller must hold pidMu. func readPidFileUnlocked(pidPath string) (*PidFileData, error) { @@ -150,12 +185,12 @@ func readPidFileUnlocked(pidPath string) (*PidFileData, error) { var data PidFileData if err := json.Unmarshal(raw, &data); err != nil { - return nil, err + return nil, fmt.Errorf("%w: %v", errInvalidPidFile, err) } // Validate PID is a positive integer. if data.PID <= 0 { - return nil, fmt.Errorf("invalid pid in pid file: %d", data.PID) + return nil, fmt.Errorf("%w: pid=%d", errInvalidPidFile, data.PID) } return &data, nil diff --git a/pkg/pid/pidfile_test.go b/pkg/pid/pidfile_test.go index 921f590ad..2da44bbbc 100644 --- a/pkg/pid/pidfile_test.go +++ b/pkg/pid/pidfile_test.go @@ -191,6 +191,22 @@ func TestReadPidFileWithCheckStalePID(t *testing.T) { } } +// TestReadPidFileWithCheckInvalidFile auto-cleans malformed PID file. +func TestReadPidFileWithCheckInvalidFile(t *testing.T) { + dir := tmpDir(t) + path := filepath.Join(dir, pidFileName) + os.WriteFile(path, []byte("not json"), 0o600) + + data := ReadPidFileWithCheck(dir) + if data != nil { + t.Error("expected nil for malformed pid file") + } + + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Error("malformed PID file should be removed") + } +} + // TestRemovePidFile removes the PID file for the current process. func TestRemovePidFile(t *testing.T) { dir := tmpDir(t) @@ -228,6 +244,40 @@ func TestRemovePidFileNonexistent(t *testing.T) { RemovePidFile(dir) } +func TestRemovePidFileIfPID(t *testing.T) { + dir := tmpDir(t) + + other := PidFileData{PID: 99999999, Token: "deadbeef12345678deadbeef12345678"} + raw, _ := json.MarshalIndent(other, "", " ") + path := filepath.Join(dir, pidFileName) + os.WriteFile(path, raw, 0o600) + + removed := RemovePidFileIfPID(dir, 99999999) + if !removed { + t.Fatal("expected RemovePidFileIfPID to remove matching pid file") + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Error("PID file should be removed for matching expected PID") + } +} + +func TestRemovePidFileIfPIDMismatch(t *testing.T) { + dir := tmpDir(t) + + other := PidFileData{PID: 99999999, Token: "deadbeef12345678deadbeef12345678"} + raw, _ := json.MarshalIndent(other, "", " ") + path := filepath.Join(dir, pidFileName) + os.WriteFile(path, raw, 0o600) + + removed := RemovePidFileIfPID(dir, 88888888) + if removed { + t.Fatal("expected RemovePidFileIfPID to keep non-matching pid file") + } + if _, err := os.Stat(path); os.IsNotExist(err) { + t.Error("PID file should NOT be removed for mismatching expected PID") + } +} + // TestReadPidFileUnlockedInvalidJSON returns error for malformed content. func TestReadPidFileUnlockedInvalidJSON(t *testing.T) { dir := tmpDir(t) diff --git a/pkg/pid/pidfile_unix.go b/pkg/pid/pidfile_unix.go index 5459d8370..7bc53b752 100644 --- a/pkg/pid/pidfile_unix.go +++ b/pkg/pid/pidfile_unix.go @@ -3,6 +3,7 @@ package pid import ( + "errors" "os" "syscall" ) @@ -18,5 +19,11 @@ func isProcessRunning(pid int) bool { return false } // Signal(nil) does not kill the process but checks existence on Unix. - return p.Signal(syscall.Signal(0)) == nil + err = p.Signal(syscall.Signal(0)) + if err == nil { + return true + } + var errno syscall.Errno + // EPERM means the process exists but we are not allowed to signal it. + return errors.As(err, &errno) && errno == syscall.EPERM } diff --git a/pkg/pid/pidfile_windows.go b/pkg/pid/pidfile_windows.go index 6a2cce793..6d8b79552 100644 --- a/pkg/pid/pidfile_windows.go +++ b/pkg/pid/pidfile_windows.go @@ -23,19 +23,19 @@ func isProcessRunning(pid int) bool { return false } - handle, _, err := procOpenProcess.Call( + handle, _, _ := procOpenProcess.Call( uintptr(processQueryLimitedInformation), 0, uintptr(pid), ) - if handle == 0 || err != nil { + if handle == 0 { return false } defer procCloseHandle.Call(handle) var exitCode uint32 - ret, _, err := procGetExitCodeProcess.Call(handle, uintptr(unsafe.Pointer(&exitCode))) - if ret == 0 || err != nil { + ret, _, _ := procGetExitCodeProcess.Call(handle, uintptr(unsafe.Pointer(&exitCode))) + if ret == 0 { return false } return exitCode == stillActive diff --git a/pkg/providers/claude_cli_provider.go b/pkg/providers/claude_cli_provider.go index 40b581490..c3d98c555 100644 --- a/pkg/providers/claude_cli_provider.go +++ b/pkg/providers/claude_cli_provider.go @@ -7,6 +7,8 @@ import ( "fmt" "os/exec" "strings" + + "github.com/sipeed/picoclaw/pkg/isolation" ) // ClaudeCliProvider implements LLMProvider using the claude CLI as a subprocess. @@ -49,7 +51,9 @@ func (p *ClaudeCliProvider) Chat( cmd.Stdout = &stdout cmd.Stderr = &stderr - if err := cmd.Run(); err != nil { + // Execute the CLI through the shared isolation wrapper so external provider + // processes honor the configured isolation policy. + if err := isolation.Run(cmd); err != nil { stderrStr := strings.TrimSpace(stderr.String()) stdoutStr := strings.TrimSpace(stdout.String()) switch { diff --git a/pkg/providers/codex_cli_provider.go b/pkg/providers/codex_cli_provider.go index 13f53ad9e..a9c8b692a 100644 --- a/pkg/providers/codex_cli_provider.go +++ b/pkg/providers/codex_cli_provider.go @@ -8,6 +8,8 @@ import ( "fmt" "os/exec" "strings" + + "github.com/sipeed/picoclaw/pkg/isolation" ) // CodexCliProvider implements LLMProvider by wrapping the codex CLI as a subprocess. @@ -56,7 +58,9 @@ func (p *CodexCliProvider) Chat( cmd.Stdout = &stdout cmd.Stderr = &stderr - err := cmd.Run() + // Execute the CLI through the shared isolation wrapper so external provider + // processes honor the configured isolation policy. + err := isolation.Run(cmd) // Parse JSONL from stdout even if exit code is non-zero, // because codex writes diagnostic noise to stderr (e.g. rollout errors) diff --git a/pkg/providers/factory_provider.go b/pkg/providers/factory_provider.go index ab7277fae..f13dc646c 100644 --- a/pkg/providers/factory_provider.go +++ b/pkg/providers/factory_provider.go @@ -160,6 +160,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err userAgent, cfg.RequestTimeout, cfg.ExtraBody, + cfg.CustomHeaders, ), modelID, nil case "azure", "azure-openai": @@ -238,6 +239,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err userAgent, cfg.RequestTimeout, cfg.ExtraBody, + cfg.CustomHeaders, ), modelID, nil case "minimax": @@ -264,6 +266,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err userAgent, cfg.RequestTimeout, extraBody, + cfg.CustomHeaders, ), modelID, nil case "anthropic": @@ -291,6 +294,7 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err userAgent, cfg.RequestTimeout, cfg.ExtraBody, + cfg.CustomHeaders, ), modelID, nil case "anthropic-messages": diff --git a/pkg/providers/factory_provider_test.go b/pkg/providers/factory_provider_test.go index b4f672f7a..c362463ae 100644 --- a/pkg/providers/factory_provider_test.go +++ b/pkg/providers/factory_provider_test.go @@ -846,6 +846,49 @@ func TestCreateProviderFromConfig_MinimaxPreservesUserExtraBody(t *testing.T) { } } +func TestCreateProviderFromConfig_CustomHeaders(t *testing.T) { + var gotSource, gotAuth string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSource = r.Header.Get("X-Source") + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`)) + })) + defer server.Close() + + cfg := &config.ModelConfig{ + ModelName: "test-headers", + Model: "openai/gpt-4o", + APIBase: server.URL, + CustomHeaders: map[string]string{"X-Source": "coding-plan", "Authorization": "Token config-auth"}, + } + cfg.SetAPIKey("test-key") + + provider, modelID, err := CreateProviderFromConfig(cfg) + if err != nil { + t.Fatalf("CreateProviderFromConfig() error = %v", err) + } + + _, err = provider.Chat( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + modelID, + nil, + ) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if gotSource != "coding-plan" { + t.Fatalf("X-Source = %q, want %q", gotSource, "coding-plan") + } + if gotAuth != "Token config-auth" { + t.Fatalf("Authorization = %q, want %q", gotAuth, "Token config-auth") + } +} + // openaiCompatResponse is the JSON response used by OpenAI-compatible providers. const openaiCompatResponse = `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}` diff --git a/pkg/providers/http_provider.go b/pkg/providers/http_provider.go index dae730536..ac91f15f6 100644 --- a/pkg/providers/http_provider.go +++ b/pkg/providers/http_provider.go @@ -24,13 +24,14 @@ func NewHTTPProvider(apiKey, apiBase, proxy string) *HTTPProvider { } func NewHTTPProviderWithMaxTokensField(apiKey, apiBase, proxy, maxTokensField string) *HTTPProvider { - return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, "", 0, nil) + return NewHTTPProviderWithMaxTokensFieldAndRequestTimeout(apiKey, apiBase, proxy, maxTokensField, "", 0, nil, nil) } func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout( apiKey, apiBase, proxy, maxTokensField, userAgent string, requestTimeoutSeconds int, extraBody map[string]any, + customHeaders map[string]string, ) *HTTPProvider { return &HTTPProvider{ delegate: openai_compat.NewProvider( @@ -40,6 +41,7 @@ func NewHTTPProviderWithMaxTokensFieldAndRequestTimeout( openai_compat.WithMaxTokensField(maxTokensField), openai_compat.WithRequestTimeout(time.Duration(requestTimeoutSeconds)*time.Second), openai_compat.WithExtraBody(extraBody), + openai_compat.WithCustomHeaders(customHeaders), openai_compat.WithUserAgent(userAgent), ), } diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 7cda033ad..d25a0fce4 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -36,6 +36,7 @@ type Provider struct { maxTokensField string // Field name for max tokens (e.g., "max_completion_tokens" for o1/glm models) httpClient *http.Client extraBody map[string]any // Additional fields to inject into request body + customHeaders map[string]string userAgent string } @@ -87,6 +88,12 @@ func WithExtraBody(extraBody map[string]any) Option { } } +func WithCustomHeaders(customHeaders map[string]string) Option { + return func(p *Provider) { + p.customHeaders = customHeaders + } +} + func NewProvider(apiKey, apiBase, proxy string, opts ...Option) *Provider { p := &Provider{ apiKey: apiKey, @@ -181,6 +188,15 @@ func (p *Provider) buildRequestBody( return requestBody } +func (p *Provider) applyCustomHeaders(req *http.Request) { + for k, v := range p.customHeaders { + if strings.TrimSpace(k) == "" { + continue + } + req.Header.Set(k, v) + } +} + func (p *Provider) Chat( ctx context.Context, messages []Message, @@ -211,6 +227,7 @@ func (p *Provider) Chat( if p.apiKey != "" { req.Header.Set("Authorization", "Bearer "+p.apiKey) } + p.applyCustomHeaders(req) resp, err := p.httpClient.Do(req) if err != nil { @@ -254,9 +271,13 @@ func (p *Provider) ChatStream( req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "text/event-stream") + if p.userAgent != "" { + req.Header.Set("User-Agent", p.userAgent) + } if p.apiKey != "" { req.Header.Set("Authorization", "Bearer "+p.apiKey) } + p.applyCustomHeaders(req) // Use a client without Timeout for streaming — the http.Client.Timeout covers // the entire request lifecycle including body reads, which would kill long streams. diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index 30aa76eb3..d140d63d6 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -710,6 +710,111 @@ func TestProviderChat_ExtraBodyOverridesOptions(t *testing.T) { } } +func TestProviderChat_CustomHeadersInjected(t *testing.T) { + var gotSource, gotAuth, gotUserAgent string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSource = r.Header.Get("X-Source") + gotAuth = r.Header.Get("Authorization") + gotUserAgent = r.Header.Get("User-Agent") + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "ok"}, + "finish_reason": "stop", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider( + "key", + server.URL, + "", + WithUserAgent("PicoClaw/Test"), + WithCustomHeaders(map[string]string{ + "X-Source": "coding-plan", + "Authorization": "Token custom-auth", + "User-Agent": "Custom-UA/1.0", + }), + ) + + _, err := p.Chat( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "gpt-4o", + nil, + ) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if gotSource != "coding-plan" { + t.Fatalf("X-Source = %q, want %q", gotSource, "coding-plan") + } + if gotAuth != "Token custom-auth" { + t.Fatalf("Authorization = %q, want %q", gotAuth, "Token custom-auth") + } + if gotUserAgent != "Custom-UA/1.0" { + t.Fatalf("User-Agent = %q, want %q", gotUserAgent, "Custom-UA/1.0") + } +} + +func TestProviderChatStream_CustomHeadersInjected(t *testing.T) { + var gotSource, gotAuth, gotUserAgent string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotSource = r.Header.Get("X-Source") + gotAuth = r.Header.Get("Authorization") + gotUserAgent = r.Header.Get("User-Agent") + + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer server.Close() + + p := NewProvider( + "key", + server.URL, + "", + WithUserAgent("PicoClaw/Test"), + WithCustomHeaders(map[string]string{ + "X-Source": "coding-plan", + "Authorization": "Token stream-auth", + "User-Agent": "Custom-UA/Stream", + }), + ) + + out, err := p.ChatStream( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "gpt-4o", + nil, + nil, + ) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + if out.Content != "ok" { + t.Fatalf("Content = %q, want %q", out.Content, "ok") + } + if gotSource != "coding-plan" { + t.Fatalf("X-Source = %q, want %q", gotSource, "coding-plan") + } + if gotAuth != "Token stream-auth" { + t.Fatalf("Authorization = %q, want %q", gotAuth, "Token stream-auth") + } + if gotUserAgent != "Custom-UA/Stream" { + t.Fatalf("User-Agent = %q, want %q", gotUserAgent, "Custom-UA/Stream") + } +} + type roundTripperFunc func(*http.Request) (*http.Response, error) func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { diff --git a/pkg/seahorse/schema_test.go b/pkg/seahorse/schema_test.go index 17879f66c..e11e6e96e 100644 --- a/pkg/seahorse/schema_test.go +++ b/pkg/seahorse/schema_test.go @@ -2,14 +2,26 @@ package seahorse import ( "database/sql" + "fmt" + "strings" + "sync/atomic" "testing" _ "modernc.org/sqlite" ) +var testDBCounter uint64 + func openTestDB(t *testing.T) *sql.DB { t.Helper() - db, err := sql.Open("sqlite", ":memory:") + + n := atomic.AddUint64(&testDBCounter, 1) + testName := strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + // Use a shared in-memory database so concurrent goroutines/connections in tests + // observe the same schema/data. + dsn := fmt.Sprintf("file:seahorse_test_%s_%d?mode=memory&cache=shared", testName, n) + + db, err := sql.Open("sqlite", dsn) if err != nil { t.Fatalf("open test db: %v", err) } diff --git a/pkg/seahorse/short_retrieval.go b/pkg/seahorse/short_retrieval.go index f7d6bf691..3e94eec14 100644 --- a/pkg/seahorse/short_retrieval.go +++ b/pkg/seahorse/short_retrieval.go @@ -68,8 +68,8 @@ type GrepSummaryResult struct { Depth int `json:"depth"` Kind SummaryKind `json:"kind"` ConversationID int64 `json:"conversationId"` - // Rank is the bm25 relevance score (negative value, closer to 0 = better match). - // Examples: -0.5 = excellent match, -2.0 = good match, -10.0 = partial match. + // Rank is the bm25 relevance score (negative value, lower = better match). + // Examples: -5.0 = excellent match, -2.0 = good match, -0.5 = partial match. Rank float64 `json:"rank,omitempty"` } @@ -79,7 +79,7 @@ type GrepMessageResult struct { Snippet string `json:"snippet"` Role string `json:"role"` ConversationID int64 `json:"conversationId"` - Rank float64 `json:"rank,omitempty"` // Relevance score (lower = better match) + Rank float64 `json:"rank,omitempty"` // Relevance score (more negative = better match) } // ExpandMessagesResult contains expanded messages. diff --git a/pkg/seahorse/tool_grep.go b/pkg/seahorse/tool_grep.go index 6502fc5c3..9671d2a7f 100644 --- a/pkg/seahorse/tool_grep.go +++ b/pkg/seahorse/tool_grep.go @@ -56,8 +56,8 @@ Returns: "hint": "No matches. Try: %keyword% for fuzzy search" } -Rank field (FTS5 mode only): bm25 relevance score, negative value where closer to 0 = better match. -Examples: -0.5=excellent, -2=good, -5=partial, -10=weak. LIKE mode (%pattern%) has no rank. +Rank field (FTS5 mode only): bm25 relevance score, negative value where more negative = higher relevance. +Examples: -5=excellent, -2=good, -0.5=partial. LIKE mode (%pattern%) has no rank. Examples: {"pattern": "authentication"} diff --git a/pkg/tools/message.go b/pkg/tools/message.go index 064065a38..5a384b37e 100644 --- a/pkg/tools/message.go +++ b/pkg/tools/message.go @@ -3,14 +3,21 @@ package tools import ( "context" "fmt" - "sync/atomic" + "sync" ) type SendCallback func(channel, chatID, content, replyToMessageID string) error +// sentTarget records the channel+chatID that the message tool sent to. +type sentTarget struct { + Channel string + ChatID string +} + type MessageTool struct { sendCallback SendCallback - sentInRound atomic.Bool // Tracks whether a message was sent in the current processing round + mu sync.Mutex + sentTargets []sentTarget // Tracks all targets sent to in the current round } func NewMessageTool() *MessageTool { @@ -53,12 +60,30 @@ func (t *MessageTool) Parameters() map[string]any { // ResetSentInRound resets the per-round send tracker. // Called by the agent loop at the start of each inbound message processing round. func (t *MessageTool) ResetSentInRound() { - t.sentInRound.Store(false) + t.mu.Lock() + t.sentTargets = t.sentTargets[:0] + t.mu.Unlock() } // HasSentInRound returns true if the message tool sent a message during the current round. func (t *MessageTool) HasSentInRound() bool { - return t.sentInRound.Load() + t.mu.Lock() + defer t.mu.Unlock() + return len(t.sentTargets) > 0 +} + +// HasSentTo returns true if the message tool sent to the specific channel+chatID +// during the current round. Used by PublishResponseIfNeeded to avoid suppressing +// the final response when the message tool only sent to a different conversation. +func (t *MessageTool) HasSentTo(channel, chatID string) bool { + t.mu.Lock() + defer t.mu.Unlock() + for _, st := range t.sentTargets { + if st.Channel == channel && st.ChatID == chatID { + return true + } + } + return false } func (t *MessageTool) SetSendCallback(callback SendCallback) { @@ -98,7 +123,10 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolRes } } - t.sentInRound.Store(true) + t.mu.Lock() + t.sentTargets = append(t.sentTargets, sentTarget{Channel: channel, ChatID: chatID}) + t.mu.Unlock() + // Silent: user already received the message directly return &ToolResult{ ForLLM: fmt.Sprintf("Message sent to %s:%s", channel, chatID), diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index d2971f3f8..a570ac9ec 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -20,6 +20,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" + "github.com/sipeed/picoclaw/pkg/isolation" ) var ( @@ -120,7 +121,7 @@ func NewExecTool(workingDir string, restrict bool, allowPaths ...[]*regexp.Regex func NewExecToolWithConfig( workingDir string, restrict bool, - config *config.Config, + cfg *config.Config, allowPaths ...[]*regexp.Regexp, ) (*ExecTool, error) { denyPatterns := make([]*regexp.Regexp, 0) @@ -131,8 +132,8 @@ func NewExecToolWithConfig( allowedPathPatterns = allowPaths[0] } - if config != nil { - execConfig := config.Tools.Exec + if cfg != nil { + execConfig := cfg.Tools.Exec enableDenyPatterns := execConfig.EnableDenyPatterns allowRemote = execConfig.AllowRemote if enableDenyPatterns { @@ -163,8 +164,8 @@ func NewExecToolWithConfig( } var timeout time.Duration - if config != nil && config.Tools.Exec.TimeoutSeconds > 0 { - timeout = time.Duration(config.Tools.Exec.TimeoutSeconds) * time.Second + if cfg != nil && cfg.Tools.Exec.TimeoutSeconds > 0 { + timeout = time.Duration(cfg.Tools.Exec.TimeoutSeconds) * time.Second } return &ExecTool{ @@ -378,7 +379,9 @@ func (t *ExecTool) runSync(ctx context.Context, command, cwd string) *ToolResult cmd.Stdout = &stdout cmd.Stderr = &stderr - if err := cmd.Start(); err != nil { + // Route shell execution through the shared isolation entry point so exec tool + // subprocesses receive the same isolation policy as other integrations. + if err := isolation.Start(cmd); err != nil { return ErrorResult(fmt.Sprintf("failed to start command: %v", err)) } @@ -521,7 +524,9 @@ func (t *ExecTool) runBackground(ctx context.Context, command, cwd string, ptyEn session.stdinWriter = stdinWriter } - if err := cmd.Start(); err != nil { + // Background sessions use the same startup path so isolation stays consistent + // with synchronous exec runs. + if err := isolation.Start(cmd); err != nil { if session.ptyMaster != nil { session.ptyMaster.Close() } diff --git a/web/backend/api/auth.go b/web/backend/api/auth.go index 22f7ec2c2..0790a6b76 100644 --- a/web/backend/api/auth.go +++ b/web/backend/api/auth.go @@ -1,8 +1,10 @@ package api import ( + "context" "crypto/subtle" "encoding/json" + "fmt" "io" "net/http" "strings" @@ -10,34 +12,47 @@ import ( "github.com/sipeed/picoclaw/web/backend/middleware" ) -// LauncherAuthRouteOpts configures dashboard token login handlers. +// PasswordStore is the interface for bcrypt-backed dashboard password persistence. +// Implemented by dashboardauth.Store; a nil value falls back to the legacy +// static-token comparison. +type PasswordStore interface { + IsInitialized(ctx context.Context) (bool, error) + SetPassword(ctx context.Context, plain string) error + VerifyPassword(ctx context.Context, plain string) (bool, error) +} + +// LauncherAuthRouteOpts configures dashboard auth handlers. type LauncherAuthRouteOpts struct { + // DashboardToken is the fallback plaintext token used when PasswordStore is + // nil or not yet initialized (env-var / config-file source, and ?token= auto-login). DashboardToken string SessionCookie string SecureCookie func(*http.Request) bool - // TokenHelp is returned on unauthenticated /api/auth/status responses (no secrets). - TokenHelp LauncherAuthTokenHelp -} - -// LauncherAuthTokenHelp tells the login UI where users can find the dashboard token. -type LauncherAuthTokenHelp struct { - EnvVarName string `json:"env_var_name"` - LogFileAbs string `json:"log_file,omitempty"` - ConfigFileAbs string `json:"config_file,omitempty"` - TrayCopyMenu bool `json:"tray_copy_menu"` - ConsoleStdout bool `json:"console_stdout"` + // PasswordStore enables bcrypt-backed password persistence. When non-nil and + // initialized, web-form login verifies against the stored hash instead of + // the plaintext DashboardToken. + PasswordStore PasswordStore + // StoreError holds the error returned when opening the password store. When + // non-nil and PasswordStore is nil, the auth endpoints surface a recovery + // message instead of an opaque 501/503. + StoreError error } type launcherAuthLoginBody struct { - Token string `json:"token"` + Password string `json:"password"` +} + +type launcherAuthSetupBody struct { + Password string `json:"password"` + Confirm string `json:"confirm"` } type launcherAuthStatusResponse struct { - Authenticated bool `json:"authenticated"` - TokenHelp *LauncherAuthTokenHelp `json:"token_help,omitempty"` + Authenticated bool `json:"authenticated"` + Initialized bool `json:"initialized"` } -// RegisterLauncherAuthRoutes registers /api/auth/login|logout|status. +// RegisterLauncherAuthRoutes registers /api/auth/login|logout|status|setup. func RegisterLauncherAuthRoutes(mux *http.ServeMux, opts LauncherAuthRouteOpts) { secure := opts.SecureCookie if secure == nil { @@ -47,22 +62,44 @@ func RegisterLauncherAuthRoutes(mux *http.ServeMux, opts LauncherAuthRouteOpts) token: opts.DashboardToken, sessionCookie: opts.SessionCookie, secureCookie: secure, - tokenHelp: opts.TokenHelp, + store: opts.PasswordStore, + storeErr: opts.StoreError, loginLimit: newLoginRateLimiter(), } mux.HandleFunc("POST /api/auth/login", h.handleLogin) mux.HandleFunc("POST /api/auth/logout", h.handleLogout) mux.HandleFunc("GET /api/auth/status", h.handleStatus) + mux.HandleFunc("POST /api/auth/setup", h.handleSetup) } type launcherAuthHandlers struct { token string sessionCookie string secureCookie func(*http.Request) bool - tokenHelp LauncherAuthTokenHelp + store PasswordStore + storeErr error // set when the store failed to open; drives recovery messages loginLimit *loginRateLimiter } +// isStoreInitialized safely queries the store. +// Returns (false, nil) when no store is configured (storeErr also nil). +// Returns (false, err) on store errors — callers must treat this as a 5xx, not as +// "uninitialized", to keep auth fail-closed. +// Exception: handleLogin swallows storeErr and falls back to token auth so +// that a corrupt DB does not lock out all access. +func (h *launcherAuthHandlers) isStoreInitialized(ctx context.Context) (bool, error) { + if h.store == nil { + if h.storeErr != nil { + return false, fmt.Errorf( + "password store unavailable (%w); "+ + "to recover, stop the application, delete the database file and restart ", + h.storeErr) + } + return false, nil + } + return h.store.IsInitialized(ctx) +} + func (h *launcherAuthHandlers) handleLogin(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") var body launcherAuthLoginBody @@ -77,10 +114,39 @@ func (h *launcherAuthHandlers) handleLogin(w http.ResponseWriter, r *http.Reques _, _ = w.Write([]byte(`{"error":"too many login attempts"}`)) return } - in := strings.TrimSpace(body.Token) - if len(in) != len(h.token) || subtle.ConstantTimeCompare([]byte(in), []byte(h.token)) != 1 { + in := strings.TrimSpace(body.Password) + var ok bool + + initialized, initErr := h.isStoreInitialized(r.Context()) + if initErr != nil { + if h.storeErr != nil { + // Store failed to open at startup — token login remains available. + initialized = false + } else { + w.WriteHeader(http.StatusInternalServerError) + writeErrorf(w, "%v", initErr) + return + } + } + + if initialized { + // Bcrypt path: verify against the stored hash. + var err error + ok, err = h.store.VerifyPassword(r.Context(), in) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + writeErrorf(w, "password verification failed: %v", err) + return + } + } else { + // Fallback: constant-time compare against the plaintext token. + ok = len(in) == len(h.token) && + subtle.ConstantTimeCompare([]byte(in), []byte(h.token)) == 1 + } + + if !ok { w.WriteHeader(http.StatusUnauthorized) - _, _ = w.Write([]byte(`{"error":"invalid token"}`)) + _, _ = w.Write([]byte(`{"error":"invalid password"}`)) return } @@ -121,23 +187,100 @@ func (h *launcherAuthHandlers) handleLogout(w http.ResponseWriter, r *http.Reque func (h *launcherAuthHandlers) handleStatus(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") - ok := false + authed := false if c, err := r.Cookie(middleware.LauncherDashboardCookieName); err == nil { - ok = subtle.ConstantTimeCompare([]byte(c.Value), []byte(h.sessionCookie)) == 1 + authed = subtle.ConstantTimeCompare([]byte(c.Value), []byte(h.sessionCookie)) == 1 } - if ok { - _, _ = w.Write([]byte(`{"authenticated":true}`)) + initialized, initErr := h.isStoreInitialized(r.Context()) + if initErr != nil { + w.WriteHeader(http.StatusServiceUnavailable) + writeErrorf(w, "%v", initErr) return } resp := launcherAuthStatusResponse{ - Authenticated: false, - TokenHelp: &h.tokenHelp, + Authenticated: authed, + Initialized: initialized, } enc, err := json.Marshal(resp) if err != nil { w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"error":"internal error"}`)) + writeErrorf(w, "marshal response failed: %v", err) return } _, _ = w.Write(enc) } + +// handleSetup sets or changes the dashboard password. +// +// Rules: +// - If the store has no password yet, the endpoint is open (no session required). +// - If a password is already set, the caller must hold a valid session cookie. +func (h *launcherAuthHandlers) handleSetup(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + if h.store == nil { + w.WriteHeader(http.StatusNotImplemented) + _, _ = w.Write([]byte(`{"error":"password store not configured"}`)) + return + } + + initialized, initErr := h.isStoreInitialized(r.Context()) + if initErr != nil { + w.WriteHeader(http.StatusServiceUnavailable) + writeErrorf(w, "%v", initErr) + return + } + + // If already initialized, require an active session (change-password flow). + if initialized { + authed := false + if c, err := r.Cookie(middleware.LauncherDashboardCookieName); err == nil { + authed = subtle.ConstantTimeCompare([]byte(c.Value), []byte(h.sessionCookie)) == 1 + } + if !authed { + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"must be authenticated to change password"}`)) + return + } + } + + var body launcherAuthSetupBody + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&body); err != nil { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"invalid JSON"}`)) + return + } + + pw := strings.TrimSpace(body.Password) + if pw == "" { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"password must not be empty"}`)) + return + } + if pw != strings.TrimSpace(body.Confirm) { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"passwords do not match"}`)) + return + } + if len([]rune(pw)) < 8 { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"password must be at least 8 characters"}`)) + return + } + + if err := h.store.SetPassword(r.Context(), pw); err != nil { + w.WriteHeader(http.StatusInternalServerError) + writeErrorf(w, "failed to save password: %v", err) + return + } + + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"status":"ok"}`)) +} + +// writeErrorf writes a JSON error response with a formatted message. +// json.Marshal is used to safely escape the message string. +func writeErrorf(w http.ResponseWriter, format string, args ...any) { + msg, _ := json.Marshal(fmt.Sprintf(format, args...)) + _, _ = w.Write([]byte(`{"error":` + string(msg) + `}`)) +} diff --git a/web/backend/api/auth_test.go b/web/backend/api/auth_test.go index d2624a440..58ffb823a 100644 --- a/web/backend/api/auth_test.go +++ b/web/backend/api/auth_test.go @@ -23,12 +23,6 @@ func TestLauncherAuthLoginAndStatus(t *testing.T) { RegisterLauncherAuthRoutes(mux, LauncherAuthRouteOpts{ DashboardToken: tok, SessionCookie: sess, - TokenHelp: LauncherAuthTokenHelp{ - EnvVarName: "PICOCLAW_LAUNCHER_TOKEN", - LogFileAbs: "/tmp/launcher.log", - TrayCopyMenu: true, - ConsoleStdout: false, - }, }) t.Run("status_unauthenticated", func(t *testing.T) { @@ -38,23 +32,20 @@ func TestLauncherAuthLoginAndStatus(t *testing.T) { t.Fatalf("status code = %d", rec.Code) } var body struct { - Authenticated bool `json:"authenticated"` - TokenHelp *LauncherAuthTokenHelp `json:"token_help"` + Authenticated bool `json:"authenticated"` + Initialized bool `json:"initialized"` } if err := json.NewDecoder(rec.Body).Decode(&body); err != nil { t.Fatal(err) } - if body.Authenticated || body.TokenHelp == nil { - t.Fatalf("unexpected body: %+v", body) - } - if body.TokenHelp.EnvVarName != "PICOCLAW_LAUNCHER_TOKEN" || body.TokenHelp.LogFileAbs != "/tmp/launcher.log" { - t.Fatalf("token_help = %+v", body.TokenHelp) + if body.Authenticated { + t.Fatalf("unexpected authenticated=true: %+v", body) } }) t.Run("login_ok", func(t *testing.T) { rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodPost, "/api/auth/login", strings.NewReader(`{"token":"`+tok+`"}`)) + req := httptest.NewRequest(http.MethodPost, "/api/auth/login", strings.NewReader(`{"password":"`+tok+`"}`)) req.Header.Set("Content-Type", "application/json") req.RemoteAddr = "127.0.0.1:12345" mux.ServeHTTP(rec, req) @@ -91,7 +82,6 @@ func TestLauncherAuthLogoutRequiresPostAndJSON(t *testing.T) { RegisterLauncherAuthRoutes(mux, LauncherAuthRouteOpts{ DashboardToken: "tok", SessionCookie: sess, - TokenHelp: LauncherAuthTokenHelp{EnvVarName: "PICOCLAW_LAUNCHER_TOKEN"}, }) rec := httptest.NewRecorder() @@ -125,11 +115,10 @@ func TestLauncherAuthLoginRateLimit(t *testing.T) { RegisterLauncherAuthRoutes(mux, LauncherAuthRouteOpts{ DashboardToken: tok, SessionCookie: sess, - TokenHelp: LauncherAuthTokenHelp{EnvVarName: "X"}, }) // 11 failing logins by wrong token; each consumes allow() slot after valid JSON. - wrongBody := `{"token":"wrong"}` + wrongBody := `{"password":"wrong"}` for i := 0; i < loginAttemptsPerIP; i++ { rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/api/auth/login", strings.NewReader(wrongBody)) @@ -187,7 +176,6 @@ func TestLauncherAuthLogoutEmptyBody(t *testing.T) { RegisterLauncherAuthRoutes(mux, LauncherAuthRouteOpts{ DashboardToken: "tok", SessionCookie: sess, - TokenHelp: LauncherAuthTokenHelp{EnvVarName: "X"}, }) rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/api/auth/logout", nil) @@ -206,7 +194,6 @@ func TestLauncherAuthLogoutRejectsTrailingJSON(t *testing.T) { RegisterLauncherAuthRoutes(mux, LauncherAuthRouteOpts{ DashboardToken: "tok", SessionCookie: sess, - TokenHelp: LauncherAuthTokenHelp{EnvVarName: "X"}, }) rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodPost, "/api/auth/logout", strings.NewReader(`{}{}`)) diff --git a/web/backend/api/gateway.go b/web/backend/api/gateway.go index b54e55bac..8994e9c60 100644 --- a/web/backend/api/gateway.go +++ b/web/backend/api/gateway.go @@ -108,6 +108,8 @@ var gatewayHealthGet = func(url string, timeout time.Duration) (*http.Response, return client.Get(url) } +var gatewayProcessMatcher = isLikelyGatewayProcess + // getGatewayHealth checks the gateway health endpoint and returns the status response. // Returns (*health.StatusResponse, statusCode, error). If error is not nil, the other values are not valid. func (h *Handler) getGatewayHealth(cfg *config.Config, timeout time.Duration) (*health.StatusResponse, int, error) { @@ -117,7 +119,7 @@ func (h *Handler) getGatewayHealth(cfg *config.Config, timeout time.Duration) (* gateway.mu.Lock() if d := gateway.pidData; d != nil && d.Port > 0 { port = d.Port - host = d.Host + host = gatewayProbeHost(d.Host) } gateway.mu.Unlock() if port == 0 { @@ -150,6 +152,150 @@ func getGatewayHealthByURL(url string, timeout time.Duration) (*health.StatusRes return &healthResponse, resp.StatusCode, nil } +// isLikelyGatewayProcess returns whether PID appears to be a picoclaw gateway +// process plus whether inspection was conclusive on this platform/environment. +func isLikelyGatewayProcess(pid int) (bool, bool) { + if pid <= 0 { + return false, true + } + + if runtime.GOOS == "windows" { + psCmd := fmt.Sprintf( + `$p=Get-CimInstance Win32_Process -Filter "ProcessId = %d"; if ($null -eq $p) { "" } else { $p.CommandLine }`, + pid, + ) + out, err := exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", psCmd).Output() + if err == nil { + cmdline := strings.TrimSpace(string(out)) + if cmdline != "" { + return looksLikeGatewayCommandLine(cmdline), true + } + } + + // Fallback: determine only whether the process still exists. + out, err = exec.Command("tasklist", "/FI", "PID eq "+strconv.Itoa(pid), "/FO", "CSV", "/NH").Output() + if err != nil { + return false, false + } + line := strings.ToLower(strings.TrimSpace(string(out))) + if line == "" { + return false, true + } + // A CSV row means the process exists, but may have a custom executable + // name we cannot classify here. + if strings.HasPrefix(line, "\"") { + if strings.Contains(line, "\"picoclaw.exe\"") { + return true, true + } + return false, false + } + if strings.Contains(line, "no tasks are running") { + return false, true + } + return false, true + } + + out, err := exec.Command("ps", "-o", "command=", "-p", strconv.Itoa(pid)).Output() + if err != nil { + return false, false + } + cmdline := strings.ToLower(strings.TrimSpace(string(out))) + if cmdline == "" { + return false, true + } + return looksLikeGatewayCommandLine(cmdline), true +} + +// looksLikeGatewayCommandLine checks whether a process command line likely +// represents "picoclaw gateway ..." regardless of executable filename. +func looksLikeGatewayCommandLine(cmdline string) bool { + fields := strings.Fields(strings.ToLower(strings.TrimSpace(cmdline))) + if len(fields) == 0 { + return false + } + for _, f := range fields { + token := strings.Trim(f, `"'`) + if token == "gateway" || strings.HasSuffix(token, "/gateway") || strings.HasSuffix(token, `\gateway`) { + return true + } + } + return false +} + +func (h *Handler) getGatewayHealthForPidData( + pidData *ppid.PidFileData, + cfg *config.Config, + timeout time.Duration, +) (*health.StatusResponse, int, error) { + if pidData == nil { + return nil, 0, errors.New("nil pid data") + } + + port := pidData.Port + if port == 0 { + port = 18790 + if cfg != nil && cfg.Gateway.Port != 0 { + port = cfg.Gateway.Port + } + } + + host := gatewayProbeHost(strings.TrimSpace(pidData.Host)) + if host == "" { + host = gatewayProbeHost(h.effectiveGatewayBindHost(cfg)) + } + if host == "" { + host = "127.0.0.1" + } + + url := "http://" + net.JoinHostPort(host, strconv.Itoa(port)) + "/health" + return getGatewayHealthByURL(url, timeout) +} + +func (h *Handler) validateGatewayPidData( + pidData *ppid.PidFileData, + cfg *config.Config, +) (ok bool, decisive bool, reason string) { + if pidData == nil || pidData.PID <= 0 { + return false, true, "invalid pid data" + } + + if gatewayProcess, inspected := gatewayProcessMatcher(pidData.PID); inspected { + if !gatewayProcess { + return false, true, "pid process command is not picoclaw gateway" + } + return true, true, "" + } + + healthResp, statusCode, err := h.getGatewayHealthForPidData(pidData, cfg, 800*time.Millisecond) + if err != nil { + return false, false, fmt.Sprintf("health probe failed: %v", err) + } + if statusCode != http.StatusOK { + return false, false, fmt.Sprintf("health endpoint returned status %d", statusCode) + } + if healthResp.PID > 0 && healthResp.PID != pidData.PID { + return false, true, fmt.Sprintf("health pid mismatch: pidFile=%d, health=%d", pidData.PID, healthResp.PID) + } + return true, true, "" +} + +func (h *Handler) sanitizeGatewayPidData(pidData *ppid.PidFileData, cfg *config.Config) *ppid.PidFileData { + if pidData == nil { + return nil + } + + ok, decisive, reason := h.validateGatewayPidData(pidData, cfg) + if ok { + return pidData + } + + logger.Warnf("ignore pid file for PID %d: %s", pidData.PID, reason) + if decisive && ppid.RemovePidFileIfPID(globalConfigDir(), pidData.PID) { + logger.Warnf("removed stale pid file for PID %d", pidData.PID) + } + return nil +} + // registerGatewayRoutes binds gateway lifecycle endpoints to the ServeMux. func (h *Handler) registerGatewayRoutes(mux *http.ServeMux) { mux.HandleFunc("GET /api/gateway/status", h.handleGatewayStatus) @@ -164,7 +310,7 @@ func (h *Handler) registerGatewayRoutes(mux *http.ServeMux) { // starts it when possible. Intended to be called by the backend at startup. func (h *Handler) TryAutoStartGateway() { // Check PID file first to detect an already-running gateway. - pidData := ppid.ReadPidFileWithCheck(globalConfigDir()) + pidData := h.sanitizeGatewayPidData(ppid.ReadPidFileWithCheck(globalConfigDir()), nil) if pidData != nil { gateway.mu.Lock() ready, reason, err := h.gatewayStartReady() @@ -357,7 +503,13 @@ func isCmdProcessAliveLocked(cmd *exec.Cmd) bool { return true } - return cmd.Process.Signal(syscall.Signal(0)) == nil + err := cmd.Process.Signal(syscall.Signal(0)) + if err == nil { + return true + } + var errno syscall.Errno + // EPERM means the process exists but cannot be signaled by this user. + return errors.As(err, &errno) && errno == syscall.EPERM } func setGatewayRuntimeStatusLocked(status string) { @@ -401,6 +553,15 @@ func gatewayStatusWithoutHealthLocked() string { return "error" } if gateway.runtimeStatus == "running" { + // For attached processes there is no waiter goroutine; degrade stale + // running state once the tracked process exits. + if !isCmdProcessAliveLocked(gateway.cmd) { + gateway.cmd = nil + gateway.owned = false + gateway.bootDefaultModel = "" + gateway.bootConfigSignature = "" + return "stopped" + } return "running" } if gateway.runtimeStatus == "error" { @@ -457,6 +618,11 @@ func stopGatewayLocked() (int, error) { } pid := gateway.cmd.Process.Pid + if !gateway.owned { + if isGateway, inspected := gatewayProcessMatcher(pid); inspected && !isGateway { + return pid, fmt.Errorf("refuse to stop non-gateway process (PID %d)", pid) + } + } // Send SIGTERM for graceful shutdown (SIGKILL on Windows) var sigErr error @@ -614,6 +780,7 @@ func (h *Handler) startGatewayLocked(initialStatus string, existingPid int) (int // Start a goroutine to probe pidFile and health, update runtime state once ready. go func() { + healthConfirmed := false for i := 0; i < 30; i++ { // try for up to 15 seconds time.Sleep(500 * time.Millisecond) gateway.mu.Lock() @@ -648,7 +815,11 @@ func (h *Handler) startGatewayLocked(initialStatus string, existingPid int) (int setGatewayRuntimeStatusLocked("running") } gateway.mu.Unlock() - return + if !healthConfirmed { + healthConfirmed = true + logger.InfoC("gateway", "Gateway health endpoint reachable; waiting for pid file") + } + continue } } }() @@ -661,7 +832,7 @@ func (h *Handler) startGatewayLocked(initialStatus string, existingPid int) (int // POST /api/gateway/start func (h *Handler) handleGatewayStart(w http.ResponseWriter, r *http.Request) { // Check PID file first to detect an already-running gateway. - pidData := ppid.ReadPidFileWithCheck(globalConfigDir()) + pidData := h.sanitizeGatewayPidData(ppid.ReadPidFileWithCheck(globalConfigDir()), nil) if pidData != nil { pid := pidData.PID gateway.mu.Lock() @@ -787,9 +958,22 @@ func (h *Handler) RestartGateway() (int, error) { gateway.mu.Lock() previousCmd := gateway.cmd + previousOwned := gateway.owned setGatewayRuntimeStatusLocked("restarting") gateway.mu.Unlock() + if previousCmd != nil && previousCmd.Process != nil && !previousOwned { + if isGateway, inspected := gatewayProcessMatcher(previousCmd.Process.Pid); inspected && !isGateway { + logger.Warnf("refuse restarting non-gateway process (PID: %d)", previousCmd.Process.Pid) + gateway.mu.Lock() + if gateway.cmd == previousCmd { + setGatewayRuntimeStatusLocked("running") + } + gateway.mu.Unlock() + return 0, fmt.Errorf("refuse to restart non-gateway process (PID %d)", previousCmd.Process.Pid) + } + } + if err = stopGatewayProcessForRestart(previousCmd); err != nil { gateway.mu.Lock() if gateway.cmd == previousCmd { @@ -901,7 +1085,7 @@ func (h *Handler) gatewayStatusData() map[string]any { } // Primary detection: read PID file and check if process is alive. - pidData := ppid.ReadPidFileWithCheck(globalConfigDir()) + pidData := h.sanitizeGatewayPidData(ppid.ReadPidFileWithCheck(globalConfigDir()), cfg) if pidData != nil { gateway.mu.Lock() gateway.pidData = pidData @@ -927,8 +1111,14 @@ func (h *Handler) gatewayStatusData() map[string]any { // (startGatewayLocked) already handles liveness detection via // pidFile polling and health fallback. gateway.mu.Lock() - data["gateway_status"] = gatewayStatusWithoutHealthLocked() - gateway.pidData = nil + status := gatewayStatusWithoutHealthLocked() + data["gateway_status"] = status + // Keep last known pidData while gateway is still in a transient + // running state; otherwise websocket proxy may lose auth token + // during short pid-file races. + if status == "stopped" || status == "error" { + gateway.pidData = nil + } gateway.mu.Unlock() } diff --git a/web/backend/api/gateway_test.go b/web/backend/api/gateway_test.go index 2ddb1fd8d..d300b657c 100644 --- a/web/backend/api/gateway_test.go +++ b/web/backend/api/gateway_test.go @@ -15,8 +15,6 @@ import ( "testing" "time" - "github.com/stretchr/testify/require" - "github.com/sipeed/picoclaw/pkg/auth" "github.com/sipeed/picoclaw/pkg/config" ppid "github.com/sipeed/picoclaw/pkg/pid" @@ -40,6 +38,36 @@ func startLongRunningProcess(t *testing.T) *exec.Cmd { return cmd } +func startGatewayLikeProcess(t *testing.T) *exec.Cmd { + t.Helper() + + var cmd *exec.Cmd + if runtime.GOOS == "windows" { + t.Skip("gateway-like process commandline check is not deterministic on Windows tests") + } + cmd = exec.Command("sh", "-c", "sleep 30 # picoclaw gateway") + + if err := cmd.Start(); err != nil { + t.Fatalf("Start() error = %v", err) + } + + return cmd +} + +func writeTestPidFile(t *testing.T, data ppid.PidFileData) string { + t.Helper() + + path := filepath.Join(globalConfigDir(), ".picoclaw.pid") + raw, err := json.MarshalIndent(data, "", " ") + if err != nil { + t.Fatalf("marshal pid file: %v", err) + } + if err := os.WriteFile(path, raw, 0o600); err != nil { + t.Fatalf("write pid file: %v", err) + } + return path +} + func mockGatewayHealthResponse(statusCode, pid int) *http.Response { return &http.Response{ StatusCode: statusCode, @@ -68,12 +96,14 @@ func resetGatewayTestState(t *testing.T) { t.Helper() originalHealthGet := gatewayHealthGet + originalProcessMatcher := gatewayProcessMatcher originalRestartGracePeriod := gatewayRestartGracePeriod originalRestartForceKillWindow := gatewayRestartForceKillWindow originalRestartPollInterval := gatewayRestartPollInterval t.Setenv("PICOCLAW_HOME", t.TempDir()) t.Cleanup(func() { gatewayHealthGet = originalHealthGet + gatewayProcessMatcher = originalProcessMatcher gatewayRestartGracePeriod = originalRestartGracePeriod gatewayRestartForceKillWindow = originalRestartForceKillWindow gatewayRestartPollInterval = originalRestartPollInterval @@ -105,6 +135,105 @@ func TestGatewayStartReady_NoDefaultModel(t *testing.T) { } } +func TestLooksLikeGatewayCommandLine(t *testing.T) { + cases := []struct { + name string + cmdline string + want bool + }{ + { + name: "default picoclaw gateway", + cmdline: "/usr/local/bin/picoclaw gateway -E", + want: true, + }, + { + name: "renamed binary with gateway subcommand", + cmdline: "/opt/bin/custom-claw gateway -E -d", + want: true, + }, + { + name: "standalone gateway binary path", + cmdline: "/opt/bin/gateway -E", + want: true, + }, + { + name: "non gateway process", + cmdline: "/bin/sleep 30", + want: false, + }, + { + name: "gateway substring only", + cmdline: "/opt/bin/gatewayd --serve", + want: false, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := looksLikeGatewayCommandLine(tc.cmdline) + if got != tc.want { + t.Fatalf("looksLikeGatewayCommandLine(%q) = %v, want %v", tc.cmdline, got, tc.want) + } + }) + } +} + +func TestValidateGatewayPidDataAcceptsHealthWhenMatcherInconclusive(t *testing.T) { + resetGatewayTestState(t) + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + + const testPID = 34567 + pidData := &ppid.PidFileData{ + PID: testPID, + Host: "127.0.0.1", + Port: 18790, + } + + gatewayProcessMatcher = func(int) (bool, bool) { return false, false } + gatewayHealthGet = func(string, time.Duration) (*http.Response, error) { + return mockGatewayHealthResponse(http.StatusOK, testPID), nil + } + + ok, decisive, reason := h.validateGatewayPidData(pidData, nil) + if !ok { + t.Fatalf("validateGatewayPidData() ok = false, want true (reason=%q)", reason) + } + if !decisive { + t.Fatalf("validateGatewayPidData() decisive = false, want true") + } +} + +func TestValidateGatewayPidDataRejectsHealthPidMismatchWhenMatcherInconclusive(t *testing.T) { + resetGatewayTestState(t) + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + + pidData := &ppid.PidFileData{ + PID: 34567, + Host: "127.0.0.1", + Port: 18790, + } + + gatewayProcessMatcher = func(int) (bool, bool) { return false, false } + gatewayHealthGet = func(string, time.Duration) (*http.Response, error) { + return mockGatewayHealthResponse(http.StatusOK, 99999), nil + } + + ok, decisive, reason := h.validateGatewayPidData(pidData, nil) + if ok { + t.Fatalf("validateGatewayPidData() ok = true, want false") + } + if !decisive { + t.Fatalf("validateGatewayPidData() decisive = false, want true") + } + if !strings.Contains(reason, "health pid mismatch") { + t.Fatalf("validateGatewayPidData() reason = %q, want contains %q", reason, "health pid mismatch") + } +} + func TestGatewayStartReady_InvalidDefaultModel(t *testing.T) { configPath := filepath.Join(t.TempDir(), "config.json") cfg := config.DefaultConfig() @@ -447,7 +576,7 @@ func TestGatewayStatusKeepsRunningWhenHealthProbeFailsAfterRunning(t *testing.T) } } -func TestGatewayStatusReportsRunningFromPidProbe(t *testing.T) { +func TestGatewayStatusKeepsPidDataWhileTrackedProcessAliveWhenPidFileUnavailable(t *testing.T) { resetGatewayTestState(t) configPath := filepath.Join(t.TempDir(), "config.json") @@ -463,6 +592,173 @@ func TestGatewayStatusReportsRunningFromPidProbe(t *testing.T) { _ = cmd.Wait() }) + gateway.mu.Lock() + gateway.cmd = cmd + gateway.pidData = &ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "existing-token", + } + setGatewayRuntimeStatusLocked("running") + gateway.mu.Unlock() + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK) + } + + gateway.mu.Lock() + defer gateway.mu.Unlock() + if gateway.pidData == nil { + t.Fatal("gateway.pidData was cleared while runtime status remained running") + } +} + +func TestGatewayStatusDowngradesRunningWhenTrackedProcessExitedAndPidFileMissing(t *testing.T) { + resetGatewayTestState(t) + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + cmd := startLongRunningProcess(t) + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + + gateway.mu.Lock() + gateway.cmd = cmd + gateway.pidData = &ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "stale-token", + } + setGatewayRuntimeStatusLocked("running") + gateway.mu.Unlock() + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK) + } + + var body map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + if got := body["gateway_status"]; got != "stopped" { + t.Fatalf("gateway_status = %#v, want %q", got, "stopped") + } + + gateway.mu.Lock() + defer gateway.mu.Unlock() + if gateway.pidData != nil { + t.Fatal("gateway.pidData should be cleared when tracked process has exited") + } +} + +func TestGatewayStatusIgnoresAndRemovesPidFileForNonGatewayProcess(t *testing.T) { + resetGatewayTestState(t) + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + cmd := startLongRunningProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + + pidPath := writeTestPidFile(t, ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "stale-token", + Host: "127.0.0.1", + Port: 18790, + }) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK) + } + + var body map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("unmarshal response: %v", err) + } + if got := body["gateway_status"]; got != "stopped" { + t.Fatalf("gateway_status = %#v, want %q", got, "stopped") + } + if _, err := os.Stat(pidPath); !os.IsNotExist(err) { + t.Fatal("stale pid file should be removed for non-gateway process") + } +} + +func TestGatewayStopRefusesNonGatewayAttachedProcess(t *testing.T) { + resetGatewayTestState(t) + if runtime.GOOS == "windows" { + t.Skip("commandline-based process type check is best-effort on Windows") + } + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + cmd := startLongRunningProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + + gateway.mu.Lock() + gateway.cmd = cmd + gateway.owned = false + setGatewayRuntimeStatusLocked("running") + gateway.mu.Unlock() + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/gateway/stop", nil) + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusInternalServerError { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusInternalServerError) + } + if !isCmdProcessAliveLocked(cmd) { + t.Fatal("non-gateway process should not be terminated by /api/gateway/stop") + } +} + +func TestGatewayStatusReportsRunningFromPidProbe(t *testing.T) { + resetGatewayTestState(t) + gatewayProcessMatcher = func(int) (bool, bool) { return true, true } + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + cmd := startGatewayLikeProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + gateway.mu.Lock() setGatewayRuntimeStatusLocked("stopped") gateway.mu.Unlock() @@ -471,8 +767,12 @@ func TestGatewayStatusReportsRunningFromPidProbe(t *testing.T) { return mockGatewayHealthResponse(http.StatusOK, cmd.Process.Pid), nil } - _, err := ppid.WritePidFile(globalConfigDir(), "localhost", 0) - require.NoError(t, err) + writeTestPidFile(t, ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "test-token", + Host: "127.0.0.1", + Port: 18790, + }) rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil) @@ -497,6 +797,7 @@ func TestGatewayStatusReportsRunningFromPidProbe(t *testing.T) { func TestGatewayStatusRequiresRestartAfterDefaultModelChange(t *testing.T) { resetGatewayTestState(t) + gatewayProcessMatcher = func(int) (bool, bool) { return true, true } configPath := filepath.Join(t.TempDir(), "config.json") cfg := config.DefaultConfig() @@ -515,16 +816,23 @@ func TestGatewayStatusRequiresRestartAfterDefaultModelChange(t *testing.T) { mux := http.NewServeMux() h.RegisterRoutes(mux) - process, err := os.FindProcess(os.Getpid()) - if err != nil { - t.Fatalf("FindProcess() error = %v", err) - } - _, err = ppid.WritePidFile(globalConfigDir(), "localhost", 0) - require.NoError(t, err) + cmd := startGatewayLikeProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + writeTestPidFile(t, ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "test-token", + Host: "127.0.0.1", + Port: 18790, + }) bootSignature := computeConfigSignature(cfg) gateway.mu.Lock() - gateway.cmd = &exec.Cmd{Process: process} + gateway.cmd = cmd gateway.bootDefaultModel = cfg.ModelList[0].ModelName gateway.bootConfigSignature = bootSignature setGatewayRuntimeStatusLocked("running") diff --git a/web/backend/api/models.go b/web/backend/api/models.go index e6749b56e..aa4a775eb 100644 --- a/web/backend/api/models.go +++ b/web/backend/api/models.go @@ -32,13 +32,14 @@ type modelResponse struct { Proxy string `json:"proxy,omitempty"` AuthMethod string `json:"auth_method,omitempty"` // Advanced fields - ConnectMode string `json:"connect_mode,omitempty"` - Workspace string `json:"workspace,omitempty"` - RPM int `json:"rpm,omitempty"` - MaxTokensField string `json:"max_tokens_field,omitempty"` - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` - ExtraBody map[string]any `json:"extra_body,omitempty"` + ConnectMode string `json:"connect_mode,omitempty"` + Workspace string `json:"workspace,omitempty"` + RPM int `json:"rpm,omitempty"` + MaxTokensField string `json:"max_tokens_field,omitempty"` + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` + ExtraBody map[string]any `json:"extra_body,omitempty"` + CustomHeaders map[string]string `json:"custom_headers,omitempty"` // Meta Enabled bool `json:"enabled"` Available bool `json:"available"` @@ -87,6 +88,7 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) { RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, ExtraBody: m.ExtraBody, + CustomHeaders: m.CustomHeaders, Enabled: m.Enabled, Available: modelStatuses[i].Available, Status: modelStatuses[i].Status, @@ -216,6 +218,14 @@ func (h *Handler) handleUpdateModel(w http.ResponseWriter, r *http.Request) { } else if len(mc.ExtraBody) == 0 { mc.ExtraBody = nil } + // Preserve existing CustomHeaders when omitted (nil), but clear it when + // the frontend sends an empty object {} to indicate the field should + // be removed. + if mc.CustomHeaders == nil { + mc.CustomHeaders = cfg.ModelList[idx].CustomHeaders + } else if len(mc.CustomHeaders) == 0 { + mc.CustomHeaders = nil + } cfg.ModelList[idx] = &mc.ModelConfig diff --git a/web/backend/api/models_test.go b/web/backend/api/models_test.go index e54d5b77c..e4297f679 100644 --- a/web/backend/api/models_test.go +++ b/web/backend/api/models_test.go @@ -430,6 +430,112 @@ func TestHandleAddModel_PersistsAPIKey(t *testing.T) { } } +func TestHandleAddModel_PersistsCustomHeaders(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/models", bytes.NewBufferString(`{ + "model_name":"new-model-headers", + "model":"openai/gpt-4o-mini", + "custom_headers":{"X-Source":"coding-plan","X-Agent":"openclaw"} + }`)) + req.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d, body=%s", rec.Code, http.StatusOK, rec.Body.String()) + } + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + if len(cfg.ModelList) != 2 { + t.Fatalf("len(model_list) = %d, want 2", len(cfg.ModelList)) + } + + added := cfg.ModelList[1] + if added.CustomHeaders == nil { + t.Fatal("custom_headers should not be nil") + } + if got := added.CustomHeaders["X-Source"]; got != "coding-plan" { + t.Fatalf("custom_headers[X-Source] = %q, want %q", got, "coding-plan") + } + if got := added.CustomHeaders["X-Agent"]; got != "openclaw" { + t.Fatalf("custom_headers[X-Agent] = %q, want %q", got, "openclaw") + } +} + +func TestHandleUpdateModel_CustomHeadersPreserveAndClear(t *testing.T) { + configPath, cleanup := setupOAuthTestEnv(t) + defer cleanup() + + cfg, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error = %v", err) + } + cfg.ModelList = []*config.ModelConfig{{ + ModelName: "editable", + Model: "openai/gpt-4o-mini", + APIKeys: config.SimpleSecureStrings("sk-existing"), + CustomHeaders: map[string]string{"X-Source": "coding-plan"}, + }} + err = config.SaveConfig(configPath, cfg) + if err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + h := NewHandler(configPath) + mux := http.NewServeMux() + h.RegisterRoutes(mux) + + // Omitted custom_headers should preserve existing value. + recPreserve := httptest.NewRecorder() + reqPreserve := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{ + "model_name":"editable", + "model":"openai/gpt-4o-mini" + }`)) + reqPreserve.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(recPreserve, reqPreserve) + if recPreserve.Code != http.StatusOK { + t.Fatalf("preserve status = %d, want %d, body=%s", recPreserve.Code, http.StatusOK, recPreserve.Body.String()) + } + + afterPreserve, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() after preserve error = %v", err) + } + if got := afterPreserve.ModelList[0].CustomHeaders["X-Source"]; got != "coding-plan" { + t.Fatalf("preserved custom_headers[X-Source] = %q, want %q", got, "coding-plan") + } + + // Empty object should clear custom_headers. + recClear := httptest.NewRecorder() + reqClear := httptest.NewRequest(http.MethodPut, "/api/models/0", bytes.NewBufferString(`{ + "model_name":"editable", + "model":"openai/gpt-4o-mini", + "custom_headers":{} + }`)) + reqClear.Header.Set("Content-Type", "application/json") + mux.ServeHTTP(recClear, reqClear) + if recClear.Code != http.StatusOK { + t.Fatalf("clear status = %d, want %d, body=%s", recClear.Code, http.StatusOK, recClear.Body.String()) + } + + afterClear, err := config.LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() after clear error = %v", err) + } + if afterClear.ModelList[0].CustomHeaders != nil { + t.Fatalf("custom_headers = %#v, want nil", afterClear.ModelList[0].CustomHeaders) + } +} + // TestHandleSetDefaultModel_RejectsNonexistentModel tests that setting a non-existent // model as default returns 404. This covers the case where virtual models (which are // filtered by SaveConfig) cannot be set as default. diff --git a/web/backend/api/pico.go b/web/backend/api/pico.go index c8ef47308..1d6b46d32 100644 --- a/web/backend/api/pico.go +++ b/web/backend/api/pico.go @@ -11,6 +11,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" + ppid "github.com/sipeed/picoclaw/pkg/pid" ) // registerPicoRoutes binds Pico Channel management endpoints to the ServeMux. @@ -57,9 +58,34 @@ func (h *Handler) handleWebSocketProxy() http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { gateway.mu.Lock() ensurePicoTokenCachedLocked(h.configPath) - gatewayAvailable := gateway.pidData != nil + cachedPID := gateway.pidData + trackedCmd := gateway.cmd gateway.mu.Unlock() + gatewayAvailable := false + // Prefer fresh PID file data when available. + if pidData := h.sanitizeGatewayPidData(ppid.ReadPidFileWithCheck(globalConfigDir()), nil); pidData != nil { + gateway.mu.Lock() + gateway.pidData = pidData + setGatewayRuntimeStatusLocked("running") + gatewayAvailable = true + gateway.mu.Unlock() + } else if cachedPID != nil { + // No PID file now: keep availability only while tracked process is + // still alive (covers short PID-file races at startup/restart). + if isCmdProcessAliveLocked(trackedCmd) { + gatewayAvailable = true + } else { + gateway.mu.Lock() + if gateway.cmd == trackedCmd { + gateway.pidData = nil + setGatewayRuntimeStatusLocked("stopped") + } + gatewayAvailable = gateway.pidData != nil + gateway.mu.Unlock() + } + } + if !gatewayAvailable { logger.Warnf("Gateway not available for WebSocket proxy") http.Error(w, "Gateway not available", http.StatusServiceUnavailable) diff --git a/web/backend/api/pico_test.go b/web/backend/api/pico_test.go index ee5586746..af5ba205f 100644 --- a/web/backend/api/pico_test.go +++ b/web/backend/api/pico_test.go @@ -11,6 +11,7 @@ import ( "strconv" "testing" + "github.com/sipeed/picoclaw/pkg/channels/pico" "github.com/sipeed/picoclaw/pkg/config" ppid "github.com/sipeed/picoclaw/pkg/pid" ) @@ -307,6 +308,13 @@ func TestHandlePicoSetup_Response(t *testing.T) { } func TestHandleWebSocketProxyReloadsGatewayTargetFromConfig(t *testing.T) { + origMatcher := gatewayProcessMatcher + gatewayProcessMatcher = func(int) (bool, bool) { return true, true } + t.Cleanup(func() { gatewayProcessMatcher = origMatcher }) + + home := t.TempDir() + t.Setenv("PICOCLAW_HOME", home) + configPath := filepath.Join(t.TempDir(), "config.json") h := NewHandler(configPath) handler := h.handleWebSocketProxy() @@ -335,6 +343,26 @@ func TestHandleWebSocketProxyReloadsGatewayTargetFromConfig(t *testing.T) { if err := config.SaveConfig(configPath, cfg); err != nil { t.Fatalf("SaveConfig() error = %v", err) } + cmd := startGatewayLikeProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + writeTestPidFile(t, ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "test-token", + Host: cfg.Gateway.Host, + Port: cfg.Gateway.Port, + }) + origPidData := gateway.pidData + origPicoToken := gateway.picoToken + t.Cleanup(func() { + ppid.RemovePidFile(globalConfigDir()) + gateway.pidData = origPidData + gateway.picoToken = origPicoToken + }) gateway.pidData = &ppid.PidFileData{} gateway.picoToken = "pico" @@ -378,6 +406,13 @@ func TestHandleWebSocketProxyReloadsGatewayTargetFromConfig(t *testing.T) { } func TestHandleWebSocketProxyLoadsCachedPicoTokenWhenMissing(t *testing.T) { + origMatcher := gatewayProcessMatcher + gatewayProcessMatcher = func(int) (bool, bool) { return true, true } + t.Cleanup(func() { gatewayProcessMatcher = origMatcher }) + + home := t.TempDir() + t.Setenv("PICOCLAW_HOME", home) + configPath := filepath.Join(t.TempDir(), "config.json") h := NewHandler(configPath) handler := h.handleWebSocketProxy() @@ -399,6 +434,22 @@ func TestHandleWebSocketProxyLoadsCachedPicoTokenWhenMissing(t *testing.T) { if err := config.SaveConfig(configPath, cfg); err != nil { t.Fatalf("SaveConfig() error = %v", err) } + cmd := startGatewayLikeProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + writeTestPidFile(t, ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "test-token", + Host: cfg.Gateway.Host, + Port: cfg.Gateway.Port, + }) + t.Cleanup(func() { + ppid.RemovePidFile(globalConfigDir()) + }) origPidData := gateway.pidData origPicoToken := gateway.picoToken @@ -426,6 +477,148 @@ func TestHandleWebSocketProxyLoadsCachedPicoTokenWhenMissing(t *testing.T) { } } +func TestHandleWebSocketProxyLoadsPidDataOnDemand(t *testing.T) { + origMatcher := gatewayProcessMatcher + gatewayProcessMatcher = func(int) (bool, bool) { return true, true } + t.Cleanup(func() { gatewayProcessMatcher = origMatcher }) + + home := t.TempDir() + t.Setenv("PICOCLAW_HOME", home) + + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + handler := h.handleWebSocketProxy() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/pico/ws" { + t.Fatalf("path = %q, want %q", r.URL.Path, "/pico/ws") + } + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, r.Header.Get(protocolKey)) + })) + defer server.Close() + + cfg := config.DefaultConfig() + cfg.Gateway.Host = "127.0.0.1" + cfg.Gateway.Port = mustGatewayTestPort(t, server.URL) + cfg.Channels.Pico.Enabled = true + cfg.Channels.Pico.SetToken("ui-token") + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + cmd := startGatewayLikeProcess(t) + t.Cleanup(func() { + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + }) + pidData := ppid.PidFileData{ + PID: cmd.Process.Pid, + Token: "test-token", + Host: cfg.Gateway.Host, + Port: cfg.Gateway.Port, + } + writeTestPidFile(t, pidData) + t.Cleanup(func() { + ppid.RemovePidFile(globalConfigDir()) + }) + + origPidData := gateway.pidData + origPicoToken := gateway.picoToken + origStatus := gateway.runtimeStatus + t.Cleanup(func() { + gateway.mu.Lock() + gateway.pidData = origPidData + gateway.picoToken = origPicoToken + gateway.runtimeStatus = origStatus + gateway.mu.Unlock() + }) + + gateway.mu.Lock() + gateway.pidData = nil + gateway.picoToken = "" + setGatewayRuntimeStatusLocked("stopped") + gateway.mu.Unlock() + + req := httptest.NewRequest(http.MethodGet, "/pico/ws?session_id=test-session", nil) + req.Header.Set(protocolKey, tokenPrefix+"ui-token") + rec := httptest.NewRecorder() + handler(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK) + } + + expected := tokenPrefix + pico.PicoTokenPrefix + pidData.Token + "ui-token" + if got := rec.Body.String(); got != expected { + t.Fatalf("forwarded protocol = %q, want %q", got, expected) + } + + gateway.mu.Lock() + defer gateway.mu.Unlock() + if gateway.pidData == nil { + t.Fatal("gateway.pidData should be loaded from pid file") + } + if gateway.runtimeStatus != "running" { + t.Fatalf("runtimeStatus = %q, want %q", gateway.runtimeStatus, "running") + } +} + +func TestHandleWebSocketProxyRejectsStalePidDataAfterProcessExit(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "config.json") + h := NewHandler(configPath) + handler := h.handleWebSocketProxy() + + cfg := config.DefaultConfig() + cfg.Channels.Pico.Enabled = true + cfg.Channels.Pico.SetToken("ui-token") + if err := config.SaveConfig(configPath, cfg); err != nil { + t.Fatalf("SaveConfig() error = %v", err) + } + + cmd := startLongRunningProcess(t) + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = cmd.Wait() + + origPidData := gateway.pidData + origPicoToken := gateway.picoToken + origCmd := gateway.cmd + origStatus := gateway.runtimeStatus + t.Cleanup(func() { + gateway.mu.Lock() + gateway.pidData = origPidData + gateway.picoToken = origPicoToken + gateway.cmd = origCmd + gateway.runtimeStatus = origStatus + gateway.mu.Unlock() + }) + + gateway.mu.Lock() + gateway.pidData = &ppid.PidFileData{PID: cmd.Process.Pid, Token: "stale-token"} + gateway.picoToken = "ui-token" + gateway.cmd = cmd + setGatewayRuntimeStatusLocked("running") + gateway.mu.Unlock() + + req := httptest.NewRequest(http.MethodGet, "/pico/ws?session_id=test-session", nil) + req.Header.Set(protocolKey, tokenPrefix+"ui-token") + rec := httptest.NewRecorder() + handler(rec, req) + + if rec.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusServiceUnavailable) + } + gateway.mu.Lock() + defer gateway.mu.Unlock() + if gateway.pidData != nil { + t.Fatal("gateway.pidData should be cleared after stale process exit is detected") + } +} + func mustGatewayTestPort(t *testing.T, rawURL string) int { t.Helper() diff --git a/web/backend/dashboardauth/sql.go b/web/backend/dashboardauth/sql.go new file mode 100644 index 000000000..94886072b --- /dev/null +++ b/web/backend/dashboardauth/sql.go @@ -0,0 +1,24 @@ +package dashboardauth + +const ( + // DBFilename is the SQLite database file stored under the PicoClaw home directory. + DBFilename = "launcher-auth.db" + + sqliteDriver = "sqlite" + // bcryptCost is deliberately high enough to slow brute-force attempts. + bcryptCost = 12 + + sqlCreateTable = ` + CREATE TABLE IF NOT EXISTS dashboard_credentials ( + id INTEGER PRIMARY KEY CHECK (id = 1), + bcrypt_hash TEXT NOT NULL + )` + + sqlCountCredentials = `SELECT COUNT(*) FROM dashboard_credentials WHERE id = 1` + + sqlUpsertHash = ` + INSERT INTO dashboard_credentials (id, bcrypt_hash) VALUES (1, ?) + ON CONFLICT(id) DO UPDATE SET bcrypt_hash = excluded.bcrypt_hash` + + sqlSelectHash = `SELECT bcrypt_hash FROM dashboard_credentials WHERE id = 1` +) diff --git a/web/backend/dashboardauth/store.go b/web/backend/dashboardauth/store.go new file mode 100644 index 000000000..44605ba22 --- /dev/null +++ b/web/backend/dashboardauth/store.go @@ -0,0 +1,94 @@ +// Package dashboardauth provides a bcrypt-backed SQLite store for the +// launcher dashboard password. The database contains a single row (id=1) +// with the bcrypt hash; no plaintext is ever persisted. +package dashboardauth + +import ( + "context" + "database/sql" + "errors" + "fmt" + "path/filepath" + + "golang.org/x/crypto/bcrypt" + _ "modernc.org/sqlite" // register "sqlite" driver +) + +// Store holds a handle to the SQLite database that stores the bcrypt hash. +type Store struct { + db *sql.DB + path string // absolute path to the SQLite file +} + +// New opens (or creates) the database inside dir, using the package's +// canonical filename. This is the preferred constructor for most callers. +// Any error is wrapped with the resolved path so callers get actionable output. +func New(dir string) (*Store, error) { + path := filepath.Join(dir, DBFilename) + s, err := Open(path) + if err != nil { + return nil, fmt.Errorf("open %q: %w", path, err) + } + return s, nil +} + +// Open opens (or creates) the SQLite database at path and migrates the schema. +func Open(path string) (*Store, error) { + db, err := sql.Open(sqliteDriver, path) + if err != nil { + return nil, err + } + if _, err = db.Exec(sqlCreateTable); err != nil { + _ = db.Close() + return nil, err + } + return &Store{db: db, path: path}, nil +} + +// Close releases the database handle. +func (s *Store) Close() error { return s.db.Close() } + +// DBPath returns the absolute path to the SQLite database file. +func (s *Store) DBPath() string { return s.path } + +// IsInitialized reports whether a password hash has been stored. +func (s *Store) IsInitialized(ctx context.Context) (bool, error) { + var n int + err := s.db.QueryRowContext(ctx, sqlCountCredentials).Scan(&n) + if err != nil { + return false, err + } + return n > 0, nil +} + +// SetPassword hashes plain with bcrypt (cost 12) and stores (or replaces) it. +// The plaintext is never written to disk. +func (s *Store) SetPassword(ctx context.Context, plain string) error { + if len([]rune(plain)) == 0 { + return errors.New("password must not be empty") + } + hash, err := bcrypt.GenerateFromPassword([]byte(plain), bcryptCost) + if err != nil { + return err + } + _, err = s.db.ExecContext(ctx, sqlUpsertHash, string(hash)) + return err +} + +// VerifyPassword returns true iff plain matches the stored bcrypt hash. +// Returns (false, nil) when no password has been set yet. +func (s *Store) VerifyPassword(ctx context.Context, plain string) (bool, error) { + var hash string + err := s.db.QueryRowContext(ctx, sqlSelectHash).Scan(&hash) + if errors.Is(err, sql.ErrNoRows) { + return false, nil + } + if err != nil { + return false, err + } + err = bcrypt.CompareHashAndPassword([]byte(hash), []byte(plain)) + if errors.Is(err, bcrypt.ErrMismatchedHashAndPassword) { + return false, nil + } + return err == nil, err +} diff --git a/web/backend/i18n.go b/web/backend/i18n.go index 106df8506..9cda9e5d5 100644 --- a/web/backend/i18n.go +++ b/web/backend/i18n.go @@ -24,8 +24,6 @@ const ( AppTooltip TranslationKey = "AppTooltip" MenuOpen TranslationKey = "MenuOpen" MenuOpenTooltip TranslationKey = "MenuOpenTooltip" - MenuCopyToken TranslationKey = "MenuCopyToken" - MenuCopyTokenHint TranslationKey = "MenuCopyTokenHint" MenuAbout TranslationKey = "MenuAbout" MenuAboutTooltip TranslationKey = "MenuAboutTooltip" MenuVersion TranslationKey = "MenuVersion" @@ -49,8 +47,6 @@ var translations = map[Language]map[TranslationKey]string{ AppTooltip: "%s - Web Console", MenuOpen: "Open Console", MenuOpenTooltip: "Open PicoClaw console in browser", - MenuCopyToken: "Copy dashboard token", - MenuCopyTokenHint: "Copy the current web console access token to the clipboard", MenuAbout: "About", MenuAboutTooltip: "About PicoClaw", MenuVersion: "Version: %s", @@ -68,8 +64,6 @@ var translations = map[Language]map[TranslationKey]string{ AppTooltip: "%s - Web Console", MenuOpen: "打开控制台", MenuOpenTooltip: "在浏览器中打开 PicoClaw 控制台", - MenuCopyToken: "复制控制台口令", - MenuCopyTokenHint: "将当前 Web 控制台访问口令复制到剪贴板", MenuAbout: "关于", MenuAboutTooltip: "关于 PicoClaw", MenuVersion: "版本: %s", diff --git a/web/backend/main.go b/web/backend/main.go index 5e9f3315f..d9ea3474c 100644 --- a/web/backend/main.go +++ b/web/backend/main.go @@ -27,6 +27,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/web/backend/api" + "github.com/sipeed/picoclaw/web/backend/dashboardauth" "github.com/sipeed/picoclaw/web/backend/launcherconfig" "github.com/sipeed/picoclaw/web/backend/middleware" "github.com/sipeed/picoclaw/web/backend/utils" @@ -49,8 +50,6 @@ var ( // Includes ?token= for same-machine dashboard login; keep serverAddr without secrets for other use. browserLaunchURL string apiHandler *api.Handler - // launcherDashboardTokenForClipboard is read by the system tray "copy token" action (GUI mode). - launcherDashboardTokenForClipboard string noBrowser *bool ) @@ -66,6 +65,24 @@ func dashboardTokenConfigHelpPath(source launcherconfig.DashboardTokenSource, la return launcherPath } +// maskSecret masks a secret for display. It always shows up to the first 3 +// runes. The last 4 runes are only appended when at least 5 runes remain +// hidden in the middle (i.e. string length >= 12), so an 8-char minimum +// password never exposes its tail. Strings of 3 chars or fewer are fully +// masked. +func maskSecret(s string) string { + runes := []rune(s) + n := len(runes) + const prefixLen, suffixLen, minHidden = 3, 4, 5 + if n < prefixLen+suffixLen+minHidden { + if n <= prefixLen { + return "**********" + } + return string(runes[:prefixLen]) + "**********" + } + return string(runes[:prefixLen]) + "**********" + string(runes[n-suffixLen:]) +} + func main() { port := flag.String("port", "18800", "Port to listen on") public := flag.Bool("public", false, "Listen on all interfaces (0.0.0.0) instead of localhost only") @@ -209,7 +226,15 @@ func main() { logger.Fatalf("Dashboard auth setup failed: %v", dashErr) } dashboardSessionCookie := middleware.SessionCookieValue(dashboardSigningKey, dashboardToken) - launcherDashboardTokenForClipboard = dashboardToken + + // Open the bcrypt password store (creates the DB file on first run). + authStore, authStoreErr := dashboardauth.New(picoHome) + if authStoreErr != nil { + logger.ErrorC("web", fmt.Sprintf("Warning: could not open auth store: %v", authStoreErr)) + authStore = nil + } else { + defer authStore.Close() + } // Determine listen address var addr string @@ -222,20 +247,11 @@ func main() { // Initialize Server components mux := http.NewServeMux() - tokenLogFileAbs := "" - if fileLoggingEnabled { - tokenLogFileAbs = filepath.Join(picoHome, logPath, logFile) - } api.RegisterLauncherAuthRoutes(mux, api.LauncherAuthRouteOpts{ DashboardToken: dashboardToken, SessionCookie: dashboardSessionCookie, - TokenHelp: api.LauncherAuthTokenHelp{ - EnvVarName: "PICOCLAW_LAUNCHER_TOKEN", - LogFileAbs: tokenLogFileAbs, - ConfigFileAbs: dashboardTokenConfigHelpPath(dashboardTokenSource, launcherPath), - TrayCopyMenu: trayOffersDashboardTokenCopy(), - ConsoleStdout: enableConsole, - }, + PasswordStore: authStore, + StoreError: authStoreErr, }) // API Routes (e.g. /api/status) @@ -284,23 +300,23 @@ func main() { fmt.Println() switch dashboardTokenSource { case launcherconfig.DashboardTokenSourceRandom: - fmt.Printf(" Dashboard token (this run): %s\n", dashboardToken) + fmt.Printf(" Dashboard password (this run): %s\n", maskSecret(dashboardToken)) case launcherconfig.DashboardTokenSourceEnv: - fmt.Printf(" Dashboard token: %s (from PICOCLAW_LAUNCHER_TOKEN)\n", dashboardToken) + fmt.Printf(" Dashboard password: from environment variable PICOCLAW_LAUNCHER_TOKEN\n") case launcherconfig.DashboardTokenSourceConfig: - fmt.Printf(" Dashboard token: %s (from %s)\n", dashboardToken, launcherPath) + fmt.Printf(" Dashboard password: configured in %s\n", launcherPath) } fmt.Println() } switch dashboardTokenSource { case launcherconfig.DashboardTokenSourceEnv: - logger.InfoC("web", "Dashboard token: environment PICOCLAW_LAUNCHER_TOKEN") + logger.InfoC("web", "Dashboard password: environment PICOCLAW_LAUNCHER_TOKEN") case launcherconfig.DashboardTokenSourceConfig: - logger.InfoC("web", fmt.Sprintf("Dashboard token: configured in %s", launcherPath)) + logger.InfoC("web", fmt.Sprintf("Dashboard password: configured in %s", launcherPath)) case launcherconfig.DashboardTokenSourceRandom: if !enableConsole { - logger.InfoC("web", "Dashboard token (this run): "+dashboardToken) + logger.InfoC("web", "Dashboard password (this run): "+maskSecret(dashboardToken)) } } diff --git a/web/backend/main_test.go b/web/backend/main_test.go index f69705179..82bf12b40 100644 --- a/web/backend/main_test.go +++ b/web/backend/main_test.go @@ -67,3 +67,31 @@ func TestDashboardTokenConfigHelpPath(t *testing.T) { }) } } + +func TestMaskSecret(t *testing.T) { + tests := []struct { + input string + want string + }{ + // Long token (>=12 chars): first 3 + 10 stars + last 4 + {"sdhjflsjdflksdf", "sdh**********ksdf"}, + {"abcdefghijklmnopqrstuvwxyz", "abc**********wxyz"}, + // Exactly 12 chars (3+4+5 hidden): suffix shown + {"abcdefghijkl", "abc**********ijkl"}, + // 8 chars (minimum password length): suffix NOT shown — only prefix+stars + {"abcdefgh", "abc**********"}, + // 11 chars (one below threshold): suffix NOT shown + {"abcdefghijk", "abc**********"}, + // 4..3 chars: prefix shown, no suffix + {"abcdefg", "abc**********"}, + {"abcd", "abc**********"}, + // <=3 chars: fully masked + {"abc", "**********"}, + {"", "**********"}, + } + for _, tt := range tests { + if got := maskSecret(tt.input); got != tt.want { + t.Errorf("maskSecret(%q) = %q, want %q", tt.input, got, tt.want) + } + } +} diff --git a/web/backend/middleware/launcher_dashboard_auth.go b/web/backend/middleware/launcher_dashboard_auth.go index 7e92fca22..c1c4c19c6 100644 --- a/web/backend/middleware/launcher_dashboard_auth.go +++ b/web/backend/middleware/launcher_dashboard_auth.go @@ -173,6 +173,8 @@ func isPublicLauncherDashboardPath(method, p string) bool { return method == http.MethodPost case "/api/auth/status": return method == http.MethodGet + case "/api/auth/setup": + return method == http.MethodPost } return false } @@ -183,7 +185,7 @@ func isPublicLauncherDashboardStatic(method, p string) bool { if method != http.MethodGet && method != http.MethodHead { return false } - if p == "/launcher-login" { + if p == "/launcher-login" || p == "/launcher-setup" { return true } if strings.HasPrefix(p, "/assets/") { diff --git a/web/backend/systray.go b/web/backend/systray.go index 744ea4611..9dcc025df 100644 --- a/web/backend/systray.go +++ b/web/backend/systray.go @@ -6,7 +6,6 @@ import ( "fmt" "fyne.io/systray" - "github.com/atotto/clipboard" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/web/backend/utils" @@ -24,7 +23,6 @@ func onReady() { // Create menu items mOpen := systray.AddMenuItem(T(MenuOpen), T(MenuOpenTooltip)) - mCopyTok := systray.AddMenuItem(T(MenuCopyToken), T(MenuCopyTokenHint)) mAbout := systray.AddMenuItem(T(MenuAbout), T(MenuAboutTooltip)) // Add version info under About menu @@ -52,17 +50,6 @@ func onReady() { logger.Errorf("Failed to open browser: %v", err) } - case <-mCopyTok.ClickedCh: - if launcherDashboardTokenForClipboard == "" { - logger.WarnC("web", "Dashboard token is empty; cannot copy") - continue - } - if err := clipboard.WriteAll(launcherDashboardTokenForClipboard); err != nil { - logger.Errorf("Failed to copy dashboard token: %v", err) - } else { - logger.InfoC("web", "Dashboard token copied to clipboard") - } - case <-mVersion.ClickedCh: // Version info - do nothing, just shows current version diff --git a/web/backend/tray_offers_copy.go b/web/backend/tray_offers_copy.go deleted file mode 100644 index 6b7d17412..000000000 --- a/web/backend/tray_offers_copy.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build (!darwin && !freebsd) || cgo - -package main - -func trayOffersDashboardTokenCopy() bool { return true } diff --git a/web/backend/tray_offers_copy_stub.go b/web/backend/tray_offers_copy_stub.go deleted file mode 100644 index 9312700f3..000000000 --- a/web/backend/tray_offers_copy_stub.go +++ /dev/null @@ -1,5 +0,0 @@ -//go:build (darwin || freebsd) && !cgo - -package main - -func trayOffersDashboardTokenCopy() bool { return false } diff --git a/web/frontend/src/api/http.ts b/web/frontend/src/api/http.ts index 0eb872f3f..347dd9373 100644 --- a/web/frontend/src/api/http.ts +++ b/web/frontend/src/api/http.ts @@ -1,14 +1,14 @@ -import { isLauncherLoginPathname } from "@/lib/launcher-login-path" +import { isLauncherAuthPathname } from "@/lib/launcher-login-path" -function isLauncherLoginPath(): boolean { +function isLauncherAuthPath(): boolean { if (typeof globalThis.location === "undefined") { return false } - if (isLauncherLoginPathname(globalThis.location.pathname || "/")) { + if (isLauncherAuthPathname(globalThis.location.pathname || "/")) { return true } try { - return isLauncherLoginPathname( + return isLauncherAuthPathname( new URL(globalThis.location.href).pathname || "/", ) } catch { @@ -18,7 +18,7 @@ function isLauncherLoginPath(): boolean { /** * Same-origin fetch that sends cookies; redirects to launcher login on 401 JSON responses. - * Skips redirect while already on the login page to avoid reload loops (e.g. gateway poll). + * Skips redirect while already on an auth page (login or setup) to avoid reload loops. */ export async function launcherFetch( input: RequestInfo | URL, @@ -33,7 +33,7 @@ export async function launcherFetch( if ( ct.includes("application/json") && typeof globalThis.location !== "undefined" && - !isLauncherLoginPath() + !isLauncherAuthPath() ) { globalThis.location.assign("/launcher-login") } diff --git a/web/frontend/src/api/launcher-auth.ts b/web/frontend/src/api/launcher-auth.ts index 4ca51993b..ed2e30687 100644 --- a/web/frontend/src/api/launcher-auth.ts +++ b/web/frontend/src/api/launcher-auth.ts @@ -1,30 +1,23 @@ /** - * Dashboard launcher token login. Uses plain fetch (not launcherFetch) to avoid - * redirect loops on 401 while on the login page. + * Dashboard launcher auth API. + * Uses plain fetch (not launcherFetch) to avoid redirect loops on auth pages. */ export async function postLauncherDashboardLogin( - token: string, + password: string, ): Promise { const res = await fetch("/api/auth/login", { method: "POST", headers: { "Content-Type": "application/json" }, credentials: "same-origin", - body: JSON.stringify({ token: token.trim() }), + body: JSON.stringify({ password: password.trim() }), }) return res.ok } -export type LauncherAuthTokenHelp = { - env_var_name: string - log_file?: string - config_file?: string - tray_copy_menu: boolean - console_stdout: boolean -} - export type LauncherAuthStatus = { authenticated: boolean - token_help?: LauncherAuthTokenHelp + /** true when a bcrypt password has been stored in the DB */ + initialized: boolean } export async function getLauncherAuthStatus(): Promise { @@ -47,3 +40,28 @@ export async function postLauncherDashboardLogout(): Promise { }) return res.ok } + +export type SetupResult = + | { ok: true } + | { ok: false; error: string } + +export async function postLauncherDashboardSetup( + password: string, + confirm: string, +): Promise { + const res = await fetch("/api/auth/setup", { + method: "POST", + headers: { "Content-Type": "application/json" }, + credentials: "same-origin", + body: JSON.stringify({ password: password.trim(), confirm: confirm.trim() }), + }) + if (res.ok) return { ok: true } + let msg = "Unknown error" + try { + const j = (await res.json()) as { error?: string } + if (j.error) msg = j.error + } catch { + /* ignore */ + } + return { ok: false, error: msg } +} diff --git a/web/frontend/src/api/models.ts b/web/frontend/src/api/models.ts index eb8d287dd..bfdd80d6d 100644 --- a/web/frontend/src/api/models.ts +++ b/web/frontend/src/api/models.ts @@ -19,6 +19,7 @@ export interface ModelInfo { request_timeout?: number thinking_level?: string extra_body?: Record + custom_headers?: Record // Meta available: boolean status: "available" | "unconfigured" | "unreachable" diff --git a/web/frontend/src/components/app-header.tsx b/web/frontend/src/components/app-header.tsx index fa1b5a488..798ac8ad5 100644 --- a/web/frontend/src/components/app-header.tsx +++ b/web/frontend/src/components/app-header.tsx @@ -2,6 +2,7 @@ import { IconBook, IconLanguage, IconLoader2, + IconLogout, IconMenu2, IconMoon, IconPlayerPlay, @@ -39,6 +40,7 @@ import { } from "@/components/ui/tooltip" import { useGateway } from "@/hooks/use-gateway.ts" import { useTheme } from "@/hooks/use-theme.ts" +import { postLauncherDashboardLogout } from "@/api/launcher-auth" export function AppHeader() { const { i18n, t } = useTranslation() @@ -47,10 +49,12 @@ export function AppHeader() { state: gwState, loading: gwLoading, canStart, + startReason, restartRequired, start, restart, stop, + error: gwError, } = useGateway() const isRunning = gwState === "running" @@ -65,6 +69,12 @@ export function AppHeader() { (gwState === "stopped" || gwState === "error") const [showStopDialog, setShowStopDialog] = React.useState(false) + const [showLogoutDialog, setShowLogoutDialog] = React.useState(false) + + const handleLogout = async () => { + await postLauncherDashboardLogout() + globalThis.location.assign("/launcher-login") + } const handleGatewayToggle = () => { if (gwLoading || isRestarting || isStopping || (!isRunning && !canStart)) { @@ -134,6 +144,23 @@ export function AppHeader() { + + + + {t("header.logout.tooltip")} + + {t("header.logout.description")} + + + + {t("common.cancel")} + void handleLogout()}> + {t("header.logout.confirm")} + + + + +
{restartRequired && ( @@ -171,38 +198,50 @@ export function AppHeader() { - {t("header.gateway.action.stop")} + {gwError ?? t("header.gateway.action.stop")} ) : ( - + + + {/* Wrap in span so the tooltip still fires when the button is disabled */} + + + + + {(gwError || (!canStart && startReason)) ? ( + {gwError ?? startReason} + ) : null} + )} {/* Theme Toggle */} + + + + + {t("header.logout.tooltip")} + +