diff --git a/internal/acp/gateway.go b/internal/acp/gateway.go index 1293f71..61e2cca 100644 --- a/internal/acp/gateway.go +++ b/internal/acp/gateway.go @@ -179,16 +179,6 @@ func handleGatewayDisconnect( return map[string]any{"accepted": true} } -func runtimeGatewayMethod(gatewayProvider string, method string) string { - if isOpenClawMode(gatewayProvider) { - switch strings.TrimSpace(method) { - case "session.start", "session.message": - return "chat.run" - } - } - return method -} - func ensureProductionGatewayConnected( server *Server, mode string, diff --git a/internal/acp/http_handler.go b/internal/acp/http_handler.go index 039cc55..bbbed62 100644 --- a/internal/acp/http_handler.go +++ b/internal/acp/http_handler.go @@ -2,6 +2,7 @@ package acp import ( "bytes" + "context" "encoding/json" "fmt" "io" @@ -9,6 +10,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "xworkmate-bridge/internal/shared" ) @@ -189,32 +191,27 @@ func (s *Server) handleRPCWithTransform( w.Header().Set("Connection", "keep-alive") } - flusher, _ := w.(http.Flusher) + streamWriter := newSafeSSEStream(r.Context(), w) writeNotification := func(message map[string]any) { if !stream { return } - shared.WriteSSE(w, message) - if flusher != nil { - flusher.Flush() - } + streamWriter.write(message) } + defer streamWriter.close() response, rpcErr := s.handleRequest(request, writeNotification) if request.ID == nil { if stream { - _, _ = w.Write([]byte("data: [DONE]\n\n")) + streamWriter.done() } return } if rpcErr != nil { envelope := shared.ErrorEnvelope(request.ID, rpcErr.Code, rpcErr.Message) if stream { - shared.WriteSSE(w, envelope) - _, _ = w.Write([]byte("data: [DONE]\n\n")) - if flusher != nil { - flusher.Flush() - } + streamWriter.write(envelope) + streamWriter.done() return } w.WriteHeader(http.StatusOK) @@ -222,11 +219,8 @@ func (s *Server) handleRPCWithTransform( return } if stream { - shared.WriteSSE(w, shared.ResultEnvelope(request.ID, response)) - _, _ = w.Write([]byte("data: [DONE]\n\n")) - if flusher != nil { - flusher.Flush() - } + streamWriter.write(shared.ResultEnvelope(request.ID, response)) + streamWriter.done() return } w.Header().Set("Content-Type", "application/json") @@ -234,6 +228,67 @@ func (s *Server) handleRPCWithTransform( _ = json.NewEncoder(w).Encode(shared.ResultEnvelope(request.ID, response)) } +type safeSSEStream struct { + ctx context.Context + w http.ResponseWriter + flusher http.Flusher + closed atomic.Bool + mu sync.Mutex +} + +func newSafeSSEStream(ctx context.Context, w http.ResponseWriter) *safeSSEStream { + flusher, _ := w.(http.Flusher) + return &safeSSEStream{ctx: ctx, w: w, flusher: flusher} +} + +func (s *safeSSEStream) write(payload map[string]any) bool { + return s.writeRaw(func() error { + return shared.WriteSSE(s.w, payload) + }) +} + +func (s *safeSSEStream) done() bool { + return s.writeRaw(func() error { + _, err := s.w.Write([]byte("data: [DONE]\n\n")) + return err + }) +} + +func (s *safeSSEStream) close() { + s.closed.Store(true) +} + +func (s *safeSSEStream) writeRaw(write func() error) (ok bool) { + if s == nil || s.closed.Load() { + return false + } + select { + case <-s.ctx.Done(): + s.closed.Store(true) + return false + default: + } + s.mu.Lock() + defer s.mu.Unlock() + if s.closed.Load() { + return false + } + defer func() { + if recover() != nil { + s.closed.Store(true) + ok = false + } + }() + if err := write(); err != nil { + s.closed.Store(true) + return false + } + if s.flusher != nil { + s.flusher.Flush() + } + return true +} + func forceOpenClawGatewayRequest(request shared.RPCRequest) (shared.RPCRequest, *shared.RPCError) { method := strings.TrimSpace(request.Method) switch method { diff --git a/internal/acp/orchestrator.go b/internal/acp/orchestrator.go index 2b87647..ae52043 100644 --- a/internal/acp/orchestrator.go +++ b/internal/acp/orchestrator.go @@ -136,10 +136,12 @@ func (o *SessionOrchestrator) runGateway( return nil, rpcErr } params = withResolvedGatewayProvider(params, gatewayProvider) - gatewayMethod := runtimeGatewayMethod(gatewayProvider, method) + if isOpenClawMode(gatewayProvider) && isSessionTaskMethod(method) { + return o.runOpenClawGatewayChat(ctx, params, gatewayProvider, turnID, notify) + } result := o.server.gateway.RequestByMode( gatewayProvider, - gatewayMethod, + method, params, 2*time.Minute, notify, @@ -165,6 +167,181 @@ func (o *SessionOrchestrator) runGateway( return payload, nil } +func (o *SessionOrchestrator) runOpenClawGatewayChat( + _ context.Context, + params map[string]any, + gatewayProvider string, + turnID string, + notify func(map[string]any), +) (map[string]any, *shared.RPCError) { + collector := newOpenClawChatCollector() + notifyWithCollection := func(message map[string]any) { + collector.observe(message) + if notify != nil { + notify(message) + } + } + chatParams, rpcErr := openClawChatSendParams(params, turnID) + if rpcErr != nil { + return nil, rpcErr + } + sendResult := o.server.gateway.RequestByMode( + gatewayProvider, + "chat.send", + chatParams, + 2*time.Minute, + notifyWithCollection, + ) + if !sendResult.OK { + return nil, gatewayRPCError(sendResult.Error, "openclaw chat.send failed") + } + sendPayload := shared.AsMap(sendResult.Payload) + runID := strings.TrimSpace(shared.StringArg(sendPayload, "runId", turnID)) + waitResult := o.server.gateway.RequestByMode( + gatewayProvider, + "agent.wait", + map[string]any{ + "runId": runID, + "timeoutMs": 120000, + }, + 2*time.Minute, + notifyWithCollection, + ) + if !waitResult.OK { + return nil, gatewayRPCError(waitResult.Error, "openclaw agent.wait failed") + } + waitPayload := shared.AsMap(waitResult.Payload) + output := collector.output() + if output == "" { + output = firstNonEmptyString(waitPayload, "output", "message", "summary", "assistantText", "text") + } + if output == "" { + output = "OpenClaw completed without displayable output." + } + return map[string]any{ + "success": true, + "output": output, + "message": output, + "summary": output, + "turnId": turnID, + "runId": runID, + "mode": router.ExecutionTargetGatewayChat, + "resolvedGatewayProviderId": gatewayProvider, + }, nil +} + +func isSessionTaskMethod(method string) bool { + switch strings.TrimSpace(method) { + case "session.start", "session.message": + return true + default: + return false + } +} + +func openClawChatSendParams(params map[string]any, turnID string) (map[string]any, *shared.RPCError) { + message := firstNonEmptyString(params, "taskPrompt", "prompt", "message") + if message == "" { + return nil, &shared.RPCError{Code: -32602, Message: "OPENCLAW_TASK_PROMPT_REQUIRED"} + } + sessionKey := openClawSessionKey(params, turnID) + chatParams := map[string]any{ + "sessionKey": sessionKey, + "message": message, + "idempotencyKey": turnID, + } + if attachments := shared.ListArg(params, "attachments"); len(attachments) > 0 { + chatParams["attachments"] = attachments + } + if thinking := strings.TrimSpace(shared.StringArg(params, "thinking", "")); thinking != "" { + chatParams["thinking"] = thinking + } + return chatParams, nil +} + +func openClawSessionKey(params map[string]any, turnID string) string { + for _, key := range []string{"threadId", "sessionId"} { + if value := strings.TrimSpace(shared.StringArg(params, key, "")); value != "" { + return value + } + } + if trimmed := strings.TrimSpace(turnID); trimmed != "" { + return trimmed + } + return "main" +} + +func gatewayRPCError(errorPayload map[string]any, fallback string) *shared.RPCError { + message := strings.TrimSpace(shared.StringArg(errorPayload, "message", fallback)) + if message == "" { + message = fallback + } + return &shared.RPCError{Code: -32002, Message: message} +} + +func firstNonEmptyString(values map[string]any, keys ...string) string { + for _, key := range keys { + if value := strings.TrimSpace(shared.StringArg(values, key, "")); value != "" { + return value + } + } + return "" +} + +type openClawChatCollector struct { + parts []string + final string +} + +func newOpenClawChatCollector() *openClawChatCollector { + return &openClawChatCollector{} +} + +func (c *openClawChatCollector) observe(notification map[string]any) { + if c == nil { + return + } + event := shared.AsMap(shared.AsMap(notification["params"])["event"]) + if len(event) == 0 || strings.TrimSpace(shared.StringArg(event, "event", "")) != "chat.run" { + return + } + payload := shared.AsMap(event["payload"]) + text := firstNonEmptyString(payload, "assistantText", "text", "message", "output", "summary") + if text == "" { + return + } + if isTerminalGatewayPayload(payload) { + c.final = text + return + } + c.parts = append(c.parts, text) +} + +func (c *openClawChatCollector) output() string { + if c == nil { + return "" + } + if strings.TrimSpace(c.final) != "" { + return strings.TrimSpace(c.final) + } + return strings.TrimSpace(strings.Join(c.parts, "")) +} + +func isTerminalGatewayPayload(payload map[string]any) bool { + if payload == nil { + return false + } + if value, ok := payload["terminal"].(bool); ok && value { + return true + } + switch strings.TrimSpace(strings.ToLower(shared.StringArg(payload, "state", ""))) { + case "complete", "completed", "done", "ok", "success", "failed", "error", "timeout", "timed_out", "cancelled", "canceled": + return true + default: + return false + } +} + func resolvedGatewayProviderID(params map[string]any, routing RoutingResult) string { for _, value := range []string{ routing.GatewayProviderID, diff --git a/internal/acp/routing_test.go b/internal/acp/routing_test.go index 78e646a..1ab4c98 100644 --- a/internal/acp/routing_test.go +++ b/internal/acp/routing_test.go @@ -490,11 +490,14 @@ func TestExecuteSessionTaskGatewayAutoConnectsLocalOpenClaw(t *testing.T) { if gateway.ConnectCount() != 1 { t.Fatalf("expected one automatic gateway connect, got %d", gateway.ConnectCount()) } - if gateway.ChatRunCount() != 1 { - t.Fatalf("expected one OpenClaw chat.run request, got %d", gateway.ChatRunCount()) + if gateway.ChatSendCount() != 1 { + t.Fatalf("expected one OpenClaw chat.send request, got %d", gateway.ChatSendCount()) } - if got := gateway.Methods(); len(got) != 2 || got[0] != "connect" || got[1] != "chat.run" { - t.Fatalf("expected connect then chat.run, got %#v", got) + if gateway.AgentWaitCount() != 1 { + t.Fatalf("expected one OpenClaw agent.wait request, got %d", gateway.AgentWaitCount()) + } + if got := gateway.Methods(); len(got) != 3 || got[0] != "connect" || got[1] != "chat.send" || got[2] != "agent.wait" { + t.Fatalf("expected connect, chat.send, then agent.wait, got %#v", got) } client := gateway.LastConnectClient() if got := client["id"]; got != "openclaw-macos" { @@ -505,7 +508,7 @@ func TestExecuteSessionTaskGatewayAutoConnectsLocalOpenClaw(t *testing.T) { } } -func TestExecuteSessionMessageGatewayUsesOpenClawChatRun(t *testing.T) { +func TestExecuteSessionMessageGatewayUsesOpenClawChatSend(t *testing.T) { gateway := newAcpFakeOpenClawGateway(t) defer gateway.Close() @@ -535,15 +538,18 @@ func TestExecuteSessionMessageGatewayUsesOpenClawChatRun(t *testing.T) { if got := response["output"]; got != "gateway pong" { t.Fatalf("expected gateway pong output, got %#v", response) } - if gateway.ChatRunCount() != 1 { - t.Fatalf("expected one OpenClaw chat.run request, got %d", gateway.ChatRunCount()) + if gateway.ChatSendCount() != 1 { + t.Fatalf("expected one OpenClaw chat.send request, got %d", gateway.ChatSendCount()) } - if got := gateway.Methods(); len(got) != 2 || got[0] != "connect" || got[1] != "chat.run" { - t.Fatalf("expected connect then chat.run, got %#v", got) + if gateway.AgentWaitCount() != 1 { + t.Fatalf("expected one OpenClaw agent.wait request, got %d", gateway.AgentWaitCount()) + } + if got := gateway.Methods(); len(got) != 3 || got[0] != "connect" || got[1] != "chat.send" || got[2] != "agent.wait" { + t.Fatalf("expected connect, chat.send, then agent.wait, got %#v", got) } } -func TestExecuteSessionTaskGatewaySurfacesOpenClawChatRunError(t *testing.T) { +func TestExecuteSessionTaskGatewaySurfacesOpenClawChatSendError(t *testing.T) { gateway := newAcpFakeOpenClawGateway(t) defer gateway.Close() @@ -568,13 +574,48 @@ func TestExecuteSessionTaskGatewaySurfacesOpenClawChatRunError(t *testing.T) { }, }) if rpcErr == nil { - t.Fatalf("expected OpenClaw chat.run error, got response: %#v", response) + t.Fatalf("expected OpenClaw chat.send error, got response: %#v", response) } if rpcErr.Code != -32002 || !strings.Contains(rpcErr.Message, "openclaw chat failed") { - t.Fatalf("expected surfaced chat.run failure, got %#v", rpcErr) + t.Fatalf("expected surfaced chat.send failure, got %#v", rpcErr) } - if got := gateway.Methods(); len(got) != 2 || got[0] != "connect" || got[1] != "chat.run" { - t.Fatalf("expected connect then chat.run, got %#v", got) + if got := gateway.Methods(); len(got) != 2 || got[0] != "connect" || got[1] != "chat.send" { + t.Fatalf("expected connect then chat.send, got %#v", got) + } +} + +func TestExecuteSessionTaskGatewaySurfacesOpenClawAgentWaitError(t *testing.T) { + gateway := newAcpFakeOpenClawGateway(t) + defer gateway.Close() + + t.Setenv("GATEWAY_RPC_URL", gateway.URL()) + t.Setenv("BRIDGE_AUTH_TOKEN", "bridge-token") + + server := NewServer() + response, rpcErr := server.executeSessionTask(task{ + req: shared.RPCRequest{ + Method: "session.start", + Params: map[string]any{ + "sessionId": "session-openclaw-wait-fail", + "threadId": "thread-openclaw-wait-fail", + "taskPrompt": "wait-error", + "workingDirectory": t.TempDir(), + "routing": map[string]any{ + "routingMode": "explicit", + "explicitExecutionTarget": "gateway", + "preferredGatewayProviderId": "openclaw", + }, + }, + }, + }) + if rpcErr == nil { + t.Fatalf("expected OpenClaw agent.wait error, got response: %#v", response) + } + if rpcErr.Code != -32002 || !strings.Contains(rpcErr.Message, "openclaw wait failed") { + t.Fatalf("expected surfaced agent.wait failure, got %#v", rpcErr) + } + if got := gateway.Methods(); len(got) != 3 || got[0] != "connect" || got[1] != "chat.send" || got[2] != "agent.wait" { + t.Fatalf("expected connect, chat.send, then agent.wait, got %#v", got) } } @@ -655,10 +696,12 @@ type acpFakeOpenClawGateway struct { server *http.Server listener net.Listener connectCount atomic.Int32 - chatRunCount atomic.Int32 + chatSendCount atomic.Int32 + agentWaitCount atomic.Int32 lastConnectClient atomic.Value mu sync.Mutex methods []string + runMessages map[string]string } func newAcpFakeOpenClawGateway(t *testing.T) *acpFakeOpenClawGateway { @@ -667,7 +710,7 @@ func newAcpFakeOpenClawGateway(t *testing.T) *acpFakeOpenClawGateway { if err != nil { t.Fatalf("listen fake openclaw gateway: %v", err) } - fake := &acpFakeOpenClawGateway{listener: listener} + fake := &acpFakeOpenClawGateway{listener: listener, runMessages: map[string]string{}} upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} mux := http.NewServeMux() mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { @@ -753,10 +796,10 @@ func newAcpFakeOpenClawGateway(t *testing.T) *acpFakeOpenClawGateway { }, }, }) - case "chat.run": - fake.chatRunCount.Add(1) + case "chat.send": + fake.chatSendCount.Add(1) params := shared.AsMap(frame["params"]) - if strings.TrimSpace(shared.StringArg(params, "taskPrompt", "")) == "fail" { + if strings.TrimSpace(shared.StringArg(params, "message", "")) == "fail" { _ = conn.WriteJSON(map[string]any{ "type": "res", "id": id, @@ -768,13 +811,75 @@ func newAcpFakeOpenClawGateway(t *testing.T) *acpFakeOpenClawGateway { }) continue } + runID := strings.TrimSpace(shared.StringArg(params, "idempotencyKey", "fake-run")) + fake.recordRunMessage(runID, strings.TrimSpace(shared.StringArg(params, "message", ""))) _ = conn.WriteJSON(map[string]any{ "type": "res", "id": id, "ok": true, "payload": map[string]any{ - "success": true, - "output": "gateway pong", + "runId": runID, + "status": "started", + }, + }) + case "agent.wait": + fake.agentWaitCount.Add(1) + params := shared.AsMap(frame["params"]) + runID := strings.TrimSpace(shared.StringArg(params, "runId", "fake-run")) + switch fake.runMessage(runID) { + case "wait-error": + _ = conn.WriteJSON(map[string]any{ + "type": "res", + "id": id, + "ok": false, + "error": map[string]any{ + "code": "OPENCLAW_WAIT_FAILED", + "message": "openclaw wait failed", + }, + }) + continue + case "wait-timeout": + _ = conn.WriteJSON(map[string]any{ + "type": "res", + "id": id, + "ok": false, + "error": map[string]any{ + "code": "TIMEOUT", + "message": "openclaw wait timeout", + }, + }) + continue + } + _ = conn.WriteJSON(map[string]any{ + "type": "event", + "event": "chat", + "seq": 1, + "payload": map[string]any{ + "runId": runID, + "state": "final", + "message": map[string]any{ + "role": "assistant", + "content": "gateway pong", + }, + }, + }) + _ = conn.WriteJSON(map[string]any{ + "type": "res", + "id": id, + "ok": true, + "payload": map[string]any{ + "runId": runID, + "status": "ok", + }, + }) + case "chat.run": + _ = conn.WriteJSON(map[string]any{ + "type": "res", + "id": id, + "ok": false, + "error": map[string]any{ + "code": "UNKNOWN_METHOD", + "message": "unknown method: chat.run", }, }) case "session.start": @@ -820,12 +925,28 @@ func (f *acpFakeOpenClawGateway) Methods() []string { return append([]string(nil), f.methods...) } +func (f *acpFakeOpenClawGateway) recordRunMessage(runID, message string) { + f.mu.Lock() + defer f.mu.Unlock() + f.runMessages[runID] = message +} + +func (f *acpFakeOpenClawGateway) runMessage(runID string) string { + f.mu.Lock() + defer f.mu.Unlock() + return f.runMessages[runID] +} + func (f *acpFakeOpenClawGateway) ConnectCount() int { return int(f.connectCount.Load()) } -func (f *acpFakeOpenClawGateway) ChatRunCount() int { - return int(f.chatRunCount.Load()) +func (f *acpFakeOpenClawGateway) ChatSendCount() int { + return int(f.chatSendCount.Load()) +} + +func (f *acpFakeOpenClawGateway) AgentWaitCount() int { + return int(f.agentWaitCount.Load()) } func (f *acpFakeOpenClawGateway) LastConnectClient() map[string]any { diff --git a/internal/acp/web_contract_test.go b/internal/acp/web_contract_test.go index 5fec216..70fb8c7 100644 --- a/internal/acp/web_contract_test.go +++ b/internal/acp/web_contract_test.go @@ -1,6 +1,7 @@ package acp import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -270,11 +271,56 @@ func TestHTTPHandlerGatewayOpenClawForcesGatewayRouting(t *testing.T) { if !strings.Contains(recorder.Body.String(), `"resolvedGatewayProviderId":"openclaw"`) { t.Fatalf("expected forced OpenClaw gateway result, got %q", recorder.Body.String()) } - if gateway.ChatRunCount() != 1 { - t.Fatalf("expected one OpenClaw chat.run, got %d", gateway.ChatRunCount()) + if gateway.ChatSendCount() != 1 { + t.Fatalf("expected one OpenClaw chat.send, got %d", gateway.ChatSendCount()) + } + if gateway.AgentWaitCount() != 1 { + t.Fatalf("expected one OpenClaw agent.wait, got %d", gateway.AgentWaitCount()) } } +func TestSafeSSEStreamDropsLateNotificationsAfterClose(t *testing.T) { + writer := &panicSSEWriter{header: http.Header{}} + stream := newSafeSSEStream(context.Background(), writer) + + stream.close() + if stream.write(map[string]any{"method": "xworkmate.gateway.push"}) { + t.Fatal("expected closed stream to drop late notification") + } + if writer.writes != 0 { + t.Fatalf("expected no write after close, got %d", writer.writes) + } + + openStream := newSafeSSEStream(context.Background(), writer) + writer.panicOnWrite = true + if openStream.write(map[string]any{"method": "xworkmate.gateway.push"}) { + t.Fatal("expected panic writer to be marked closed") + } + if openStream.write(map[string]any{"method": "xworkmate.gateway.push"}) { + t.Fatal("expected writes after panic to stay closed") + } +} + +type panicSSEWriter struct { + header http.Header + writes int + panicOnWrite bool +} + +func (w *panicSSEWriter) Header() http.Header { + return w.header +} + +func (w *panicSSEWriter) Write(payload []byte) (int, error) { + w.writes++ + if w.panicOnWrite { + panic("closed response writer") + } + return len(payload), nil +} + +func (w *panicSSEWriter) WriteHeader(int) {} + func TestHTTPHandlerPingRequiresBearerAuthorizationWhenBridgeAuthTokenConfigured(t *testing.T) { t.Setenv("BRIDGE_AUTH_TOKEN", "bridge-test-token") t.Setenv("BRIDGE_CONFIG_PATH", "../../example/config.yaml") diff --git a/internal/shared/rpc.go b/internal/shared/rpc.go index a4878d1..3b73656 100644 --- a/internal/shared/rpc.go +++ b/internal/shared/rpc.go @@ -40,9 +40,10 @@ func DecodeRPCRequest(payload []byte) (RPCRequest, error) { return request, nil } -func WriteSSE(w http.ResponseWriter, payload map[string]any) { +func WriteSSE(w http.ResponseWriter, payload map[string]any) error { encoded, _ := json.Marshal(payload) - _, _ = fmt.Fprintf(w, "data: %s\n\n", encoded) + _, err := fmt.Fprintf(w, "data: %s\n\n", encoded) + return err } func ResultEnvelope(id any, result map[string]any) map[string]any {