From 0ec8ce8c8c914d68a4adab5b7c58241fb8e431eb Mon Sep 17 00:00:00 2001 From: Gurney Buchanan Date: Tue, 20 Oct 2026 16:12:59 -0400 Subject: [PATCH] Drive crush through its server so runs are live, not a black box MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit crush runs previously went through `crush run`, whose quiet stdout made the monitor look frozen for minutes at a time and reported no token usage, so budget bands and spend declarations never had data and aborts were process kills. The crush adapter now implements the new SessionRunner interface: it spawns `crush server` on a private Unix socket, creates a workspace and session, holds the SSE event stream open (which is what keeps the workspace alive), and drives the prompt through the agent endpoint with a RunID. The event stream maps onto the runner's normalized events: live tool calls with their arguments, finished thinking blocks, assistant text, and cumulative token usage โ€” so the monitor shows what the agent is doing while it does it, the budget bar moves on real numbers, spend rows are declared, and run_complete (correlated by RunID, foreign ones ignored) terminates the run with its final text or error. Aborts and budget refusals now cancel the remote session through the cancel endpoint before tearing the server down. Permission and question requests cannot be satisfied unattended; they fail closed (exit 2) naming the tool, with a note pointing at pre-allowing tools in a .crush/crushrc at the harness root. When the server cannot start, the run falls back to the old `crush run` path. The API surface was verified live against a real server (v0.96.1) and the whole flow is tested against an in-process fake speaking the same v1 API: happy path, permission fail-closed with remote cancel, run_complete errors, foreign-run filtering, and model override via the config/model endpoint (scope 1, model_type large, provider resolved from the `crush models` listing). ๐Ÿ’˜ Generated with Crush Assisted-by: Crush:rits/zai-org/glm-5-3 --- CHANGELOG.md | 7 + README.md | 2 +- cmd/run.go | 2 + internal/harness/crush.go | 6 +- internal/harness/crush_serve.go | 558 +++++++++++++++++++++++++++ internal/harness/crush_serve_test.go | 299 ++++++++++++++ internal/harness/harness.go | 32 +- internal/runner/runner.go | 135 ++++--- internal/runner/session_test.go | 180 +++++++++ internal/tui/monitor.go | 18 +- 10 files changed, 1183 insertions(+), 56 deletions(-) create mode 100644 internal/harness/crush_serve.go create mode 100644 internal/harness/crush_serve_test.go create mode 100644 internal/runner/session_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 51d3f67..8a0b2c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,3 +43,10 @@ Initial implementation. browser rebuilds its table columns and re-renders its markdown at the current width; plain (non-TTY) output and summary lines wrap to the terminal width instead of a fixed 200 characters. +- crush runs now execute through `crush server` in client mode: thinking, + tool calls and cumulative token usage stream live into the monitor + (run_complete is correlated by RunID), budget bands and spend + declarations work on real usage, and aborts cancel the remote session + through the cancel endpoint before teardown. Permission and question + requests cannot be satisfied unattended and fail closed (exit 2) naming + the tool. Runs fall back to `crush run` when the server cannot start. diff --git a/README.md b/README.md index 38466ec..1821caf 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,7 @@ Antigravity (agy) is deferred: it has no headless execution path. | codex | `codex exec - --json --output-last-message` | โœ“ | โœ“ | `workspace-write` sandbox, approvals never | | pi | `pi --mode json` | โœ“ | โœ“ | tool grants mapped onto pi's builtin tools | | opencode | `opencode run --format json --auto` | partial | โ€” | `--auto` | -| crush | `crush run` (stdin) | text only | โ€” | your crushrc `permissions allow` entries | +| crush | `crush server` (client mode): workspace + SSE events, clean cancel; falls back to `crush run` | โœ“ live | โœ“ | workspace permission level; unanswered permission/question requests fail closed (exit 2) | | goose | experimental (`TRAUST_CLI_GOOSE_ARGS` overrides the template) | โ€” | โ€” | goose's own configuration | Every adapter prints its honest caveats before a run. `--unsafe-full-permission` diff --git a/cmd/run.go b/cmd/run.go index 769bdf4..b49a2c7 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -248,6 +248,8 @@ func (plainSink) Event(ev harness.Event) { switch ev.Kind { case harness.EventText: fmt.Println(tui.Wrap(ev.Text, tui.TerminalWidth())) + case harness.EventThinking: + fmt.Println(tui.Wrap(styles.Muted.Render("ยท "+ev.Text), tui.TerminalWidth())) case harness.EventToolCall: fmt.Println(styles.Accent.Render("โ†’ " + ev.Tool)) case harness.EventResult: diff --git a/internal/harness/crush.go b/internal/harness/crush.go index 72c6e79..b28d974 100644 --- a/internal/harness/crush.go +++ b/internal/harness/crush.go @@ -39,9 +39,9 @@ func (c *Crush) Command(ctx context.Context, spec *RunSpec) *exec.Cmd { func (c *Crush) Notes() []string { return []string{ - "crush run emits plain text, not a structured stream: tool calls are not distinguished live", - "token usage is not reported by crush run: the spend declaration is skipped", - "headless permission posture comes from your crushrc `permissions allow` entries โ€” pre-allow the tools the skills need", + "running through crush server mode: thinking, tool calls and token usage stream live into the monitor", + "permission and question requests cannot be answered unattended: the run fails closed (exit 2) naming the tool โ€” pre-allow the tools skills need in a .crush/crushrc at the harness root (`permissions allow`)", + "if the server cannot start, the run falls back to `crush run` (plain text, no live events, no usage)", } } diff --git a/internal/harness/crush_serve.go b/internal/harness/crush_serve.go new file mode 100644 index 0000000..23ce4b6 --- /dev/null +++ b/internal/harness/crush_serve.go @@ -0,0 +1,558 @@ +package harness + +import ( + "bufio" + "bytes" + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "time" +) + +// crushServerURL, when set, points the crush adapter at an already-running +// server (unix:///path/to/socket) instead of spawning one โ€” used by tests +// and reserved for a future persistent-daemon mode. +func crushServerURL() string { return os.Getenv("TRAUST_CLI_CRUSH_SERVER_URL") } + +// RunSession implements SessionRunner: drive the skill through `crush +// server` instead of `crush run`. The server gives the monitor everything +// the one-shot process could not: live structured events (thinking, text, +// tool calls), cumulative token usage, a clean cancel endpoint, and a +// RunID-correlated run_complete terminal signal. Workspaces live exactly as +// long as an attached SSE client, so the stream stays open for the whole +// run; permission and question requests cannot be answered unattended and +// fail closed (needs-input, exit 2). +func (c *Crush) RunSession(ctx context.Context, spec *RunSpec, emit func(Event)) (SessionResult, error) { + srv, err := startCrushServer(ctx, c.Binary(), spec) + if err != nil { + return SessionResult{}, fmt.Errorf("%w: %v", ErrServerUnsupported, err) + } + defer srv.shutdown() + + clientID := randomID() + wsID, err := srv.createWorkspace(ctx, spec.Cwd, clientID) + if err != nil { + return SessionResult{}, fmt.Errorf("create workspace: %w", err) + } + sessID, err := srv.createSession(ctx, wsID) + if err != nil { + return SessionResult{}, fmt.Errorf("create session: %w", err) + } + + // The SSE stream attaches the client; without it the workspace is torn + // down and current-session would 409 ("client not attached"). + body, cancelSSE, err := srv.openEvents(ctx, wsID, clientID) + if err != nil { + return SessionResult{}, fmt.Errorf("open event stream: %w", err) + } + defer cancelSSE() + + if err := srv.setCurrentSession(ctx, wsID, clientID, sessID); err != nil { + return SessionResult{}, fmt.Errorf("attach to session: %w", err) + } + + if spec.Model != "" { + if err := srv.overrideModel(ctx, c.Binary(), spec.Env, wsID, spec.Model); err != nil { + return SessionResult{}, fmt.Errorf("override model: %w", err) + } + } + + runID := randomID() + if err := srv.sendMessage(ctx, wsID, sessID, runID, spec.Prompt); err != nil { + return SessionResult{}, fmt.Errorf("send prompt: %w", err) + } + + return srv.readEvents(ctx, body, wsID, sessID, runID, emit) +} + +// crushServer owns one `crush server` process (or borrows an external one). +type crushServer struct { + client *http.Client + base string + cmd *exec.Cmd + sockDir string + external bool +} + +// startCrushServer spawns `crush server` on a private Unix socket and waits +// for /v1/health. Any spawn or readiness failure maps to ErrServerUnsupported +// upstream so the runner can fall back to `crush run`. +func startCrushServer(ctx context.Context, bin string, spec *RunSpec) (*crushServer, error) { + if url := crushServerURL(); url != "" { + client, base, err := unixHTTPClient(url) + if err != nil { + return nil, err + } + return &crushServer{client: client, base: base, external: true}, nil + } + + sockDir, err := os.MkdirTemp("", "traust-cli-crush-*") + if err != nil { + return nil, err + } + sock := filepath.Join(sockDir, "crush.sock") + cmd := exec.CommandContext(context.WithoutCancel(ctx), bin, "server", "--host", "unix://"+sock) + cmd.Dir = spec.Cwd + cmd.Env = spec.Env + if spec.OutDir != "" { + log, lerr := os.Create(filepath.Join(spec.OutDir, "crush-server.log")) + if lerr == nil { + cmd.Stderr = log + cmd.Stdout = log + defer log.Close() + } + } + if err := cmd.Start(); err != nil { + os.RemoveAll(sockDir) + return nil, fmt.Errorf("start %s server: %w", bin, err) + } + + client, base, err := unixHTTPClient("unix://" + sock) + if err != nil { + cmd.Process.Kill() + os.RemoveAll(sockDir) + return nil, err + } + srv := &crushServer{client: client, base: base, cmd: cmd, sockDir: sockDir} + + // Readiness: poll /v1/health. + deadline := time.Now().Add(25 * time.Second) + for time.Now().Before(deadline) { + select { + case <-ctx.Done(): + srv.shutdown() + return nil, ctx.Err() + default: + } + if srv.healthy(ctx) { + return srv, nil + } + if cmd.ProcessState != nil && cmd.ProcessState.Exited() { + os.RemoveAll(sockDir) + return nil, fmt.Errorf("crush server exited during startup") + } + time.Sleep(250 * time.Millisecond) + } + srv.shutdown() + return nil, fmt.Errorf("crush server did not become healthy within 25s") +} + +// unixHTTPClient builds an HTTP client that dials a Unix socket (or a +// normal http(s) URL when the override is one). +func unixHTTPClient(url string) (*http.Client, string, error) { + if !strings.HasPrefix(url, "unix://") { + return &http.Client{Timeout: 60 * time.Second}, url, nil + } + sock := strings.TrimPrefix(url, "unix://") + transport := &http.Transport{ + DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, "unix", sock) + }, + } + return &http.Client{Transport: transport}, "http://crush.local", nil +} + +func (s *crushServer) healthy(ctx context.Context) bool { + cctx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + req, _ := http.NewRequestWithContext(cctx, http.MethodGet, s.base+"/v1/health", nil) + resp, err := s.client.Do(req) + if err != nil { + return false + } + resp.Body.Close() + return resp.StatusCode == http.StatusOK +} + +func (s *crushServer) shutdown() { + if s.external { + return + } + // Ask nicely first so crush can flush its session store. + cctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + req, _ := http.NewRequestWithContext(cctx, http.MethodPost, s.base+"/v1/control", strings.NewReader(`{"command":"shutdown"}`)) + req.Header.Set("Content-Type", "application/json") + if resp, err := s.client.Do(req); err == nil { + resp.Body.Close() + } + if s.cmd == nil || s.cmd.Process == nil { + return + } + done := make(chan struct{}) + go func() { s.cmd.Wait(); close(done) }() + select { + case <-done: + case <-time.After(5 * time.Second): + s.cmd.Process.Kill() + <-done + } + os.RemoveAll(s.sockDir) +} + +// call is a one-shot JSON request helper. +func (s *crushServer) call(ctx context.Context, method, path string, body any, out any) (int, error) { + var reader io.Reader + if body != nil { + data, err := json.Marshal(body) + if err != nil { + return 0, err + } + reader = bytes.NewReader(data) + } + req, err := http.NewRequestWithContext(ctx, method, s.base+path, reader) + if err != nil { + return 0, err + } + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + resp, err := s.client.Do(req) + if err != nil { + return 0, err + } + defer resp.Body.Close() + data, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + if err != nil { + return resp.StatusCode, err + } + if out != nil && len(data) > 0 { + if err := json.Unmarshal(data, out); err != nil { + return resp.StatusCode, fmt.Errorf("decode response: %w", err) + } + } + return resp.StatusCode, nil +} + +func (s *crushServer) createWorkspace(ctx context.Context, path, clientID string) (string, error) { + var ws struct { + ID string `json:"id"` + } + code, err := s.call(ctx, http.MethodPost, "/v1/workspaces", map[string]string{"path": path, "client_id": clientID}, &ws) + if err != nil || code != http.StatusOK { + return "", fmt.Errorf("POST /v1/workspaces: %d %v", code, err) + } + return ws.ID, nil +} + +func (s *crushServer) createSession(ctx context.Context, wsID string) (string, error) { + var sess struct { + ID string `json:"id"` + } + code, err := s.call(ctx, http.MethodPost, "/v1/workspaces/"+wsID+"/sessions", map[string]string{}, &sess) + if err != nil || code != http.StatusOK { + return "", fmt.Errorf("POST sessions: %d %v", code, err) + } + return sess.ID, nil +} + +// openEvents starts the SSE subscription. The returned closer must be +// called to detach (which releases the workspace). +func (s *crushServer) openEvents(ctx context.Context, wsID, clientID string) (io.ReadCloser, func(), error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + s.base+"/v1/workspaces/"+wsID+"/events?client_id="+clientID, nil) + if err != nil { + return nil, nil, err + } + req.Header.Set("Accept", "text/event-stream") + resp, err := s.client.Do(req) + if err != nil { + return nil, nil, err + } + if resp.StatusCode != http.StatusOK { + resp.Body.Close() + return nil, nil, fmt.Errorf("GET events: %d", resp.StatusCode) + } + return resp.Body, func() { resp.Body.Close() }, nil +} + +func (s *crushServer) setCurrentSession(ctx context.Context, wsID, clientID, sessID string) error { + code, err := s.call(ctx, http.MethodPost, + "/v1/workspaces/"+wsID+"/current-session?client_id="+clientID, + map[string]string{"session_id": sessID}, nil) + if err != nil || code != http.StatusOK && code != http.StatusConflict { + return fmt.Errorf("POST current-session: %d %v", code, err) + } + return nil +} + +// overrideModel maps a model id onto {provider, model} via `crush models` +// and sets it as the workspace's large model. Scope marshals as an int +// (ScopeWorkspace = 1); model_type as "large". +func (s *crushServer) overrideModel(ctx context.Context, bin string, env []string, wsID, model string) error { + provider, modelID, err := resolveCrushProvider(ctx, bin, env, model) + if err != nil { + return err + } + body := map[string]any{ + "scope": 1, + "model_type": "large", + "model": map[string]string{"provider": provider, "model": modelID}, + } + if code, err := s.call(ctx, http.MethodPost, "/v1/workspaces/"+wsID+"/config/model", body, nil); err != nil || code != http.StatusOK { + return fmt.Errorf("POST config/model: %d %v", code, err) + } + if code, err := s.call(ctx, http.MethodPost, "/v1/workspaces/"+wsID+"/agent/update", nil, nil); err != nil || code != http.StatusOK { + return fmt.Errorf("POST agent/update: %d %v", code, err) + } + return nil +} + +// resolveCrushProvider resolves a bare or prefixed model id to its provider +// using the `crush models` listing ("PriceTag/rits/zai-org/glm-5-3" style: +// the provider is the FIRST segment; the model id may itself contain +// slashes). +func resolveCrushProvider(ctx context.Context, bin string, env []string, model string) (string, string, error) { + if provider, rest, ok := strings.Cut(model, "/"); ok && provider != "" && rest != "" { + // Only treat as provider/model when the whole id appears in the + // listing under that provider; otherwise fall through to listing + // resolution (model ids contain slashes more often than not). + if _, _, err := matchCrushModelListing(ctx, bin, env, provider+"/"+rest, model); err == nil { + return provider, rest, nil + } + } + line, _, err := matchCrushModelListing(ctx, bin, env, "", model) + if err != nil { + return "", "", err + } + provider, rest, _ := strings.Cut(line, "/") + return provider, rest, nil +} + +// matchCrushModelListing finds the full listing line for a model (exact or +// provider-prefixed match), running `crush models` to get the listing. +func matchCrushModelListing(ctx context.Context, bin string, env []string, wantPrefix, model string) (string, string, error) { + cctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + cmd := exec.CommandContext(cctx, bin, "models") + cmd.Env = env + out, err := cmd.Output() + if err != nil { + return "", "", fmt.Errorf("crush models: %w", err) + } + for _, line := range strings.Split(string(out), "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + if line == model || strings.HasSuffix(line, "/"+model) { + if wantPrefix != "" && !strings.HasPrefix(line, wantPrefix+"/") { + continue + } + return line, line, nil + } + } + return "", "", fmt.Errorf("model %q is not in `crush models`", model) +} + +func (s *crushServer) sendMessage(ctx context.Context, wsID, sessID, runID, prompt string) error { + code, err := s.call(ctx, http.MethodPost, "/v1/workspaces/"+wsID+"/agent", map[string]any{ + "session_id": sessID, + "run_id": runID, + "prompt": prompt, + }, nil) + if err != nil || code != http.StatusAccepted { + return fmt.Errorf("POST agent: %d %v", code, err) + } + return nil +} + +func (s *crushServer) cancelSession(ctx context.Context, wsID, sessID string) { + cctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + s.call(cctx, http.MethodPost, "/v1/workspaces/"+wsID+"/agent/sessions/"+sessID+"/cancel", nil, nil) +} + +// readEvents consumes the SSE stream until run_complete for our run (or +// context cancellation, which cancels the remote run first). +func (s *crushServer) readEvents(ctx context.Context, body io.Reader, wsID, sessID, runID string, emit func(Event)) (SessionResult, error) { + result := SessionResult{} + seenParts := map[string]int{} + + reader := bufio.NewReaderSize(body, 64*1024) + for { + if ctx.Err() != nil { + s.cancelSession(ctx, wsID, sessID) + result.Cancelled = true + return result, nil + } + line, err := reader.ReadString('\n') + if err != nil { + if ctx.Err() != nil { + s.cancelSession(ctx, wsID, sessID) + result.Cancelled = true + return result, nil + } + if errors.Is(err, io.EOF) { + return result, fmt.Errorf("event stream ended without run_complete") + } + return result, fmt.Errorf("event stream: %w", err) + } + data, ok := strings.CutPrefix(strings.TrimRight(line, "\r\n"), "data: ") + if !ok || data == "" { + continue + } + if done := s.handleSSE(data, sessID, runID, seenParts, emit, &result); done { + return result, nil + } + } +} + +// handleSSE maps one SSE data frame onto normalized events. Returns true +// when the terminal run_complete for our run arrived. +func (s *crushServer) handleSSE(data, sessID, runID string, seenParts map[string]int, emit func(Event), result *SessionResult) bool { + var envelope struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` + } + if err := json.Unmarshal([]byte(data), &envelope); err != nil { + return false + } + var inner struct { + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` + } + if err := json.Unmarshal(envelope.Payload, &inner); err != nil { + return false + } + + switch envelope.Type { + case "message": + var msg struct { + ID string `json:"id"` + Role string `json:"role"` + Parts []struct { + Type string `json:"type"` + Data json.RawMessage `json:"data"` + } `json:"parts"` + } + if err := json.Unmarshal(inner.Payload, &msg); err != nil { + return false + } + if msg.Role != "assistant" { + return false + } + for i := seenParts[msg.ID]; i < len(msg.Parts); i++ { + part := msg.Parts[i] + switch part.Type { + case "text": + var p struct { + Text string `json:"text"` + Hidden bool `json:"hidden"` + } + if json.Unmarshal(part.Data, &p) == nil && p.Text != "" && !p.Hidden { + emit(Event{Kind: EventText, Text: p.Text}) + } + case "reasoning": + // Reasoning parts are rewritten as they grow; emit only the + // finished block so the transcript stays signal, not spam. + var p struct { + Thinking string `json:"thinking"` + FinishedAt int64 `json:"finished_at"` + } + if json.Unmarshal(part.Data, &p) == nil && p.FinishedAt != 0 && p.Thinking != "" { + emit(Event{Kind: EventThinking, Text: p.Thinking}) + } + case "tool_call": + var p struct { + Name string `json:"name"` + Input string `json:"input"` + } + if json.Unmarshal(part.Data, &p) == nil && p.Name != "" { + emit(Event{Kind: EventToolCall, Tool: p.Name, Text: firstLine(p.Input, 160)}) + } + } + } + seenParts[msg.ID] = len(msg.Parts) + + case "session": + var sess struct { + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Cost float64 `json:"cost"` + } + if err := json.Unmarshal(inner.Payload, &sess); err == nil { + if sess.PromptTokens > 0 || sess.CompletionTokens > 0 || sess.Cost > 0 { + emit(Event{Kind: EventUsage, TokensIn: sess.PromptTokens, TokensOut: sess.CompletionTokens, CostUSD: sess.Cost}) + } + } + + case "permission_request", "question_batch_request": + var req struct { + ToolName string `json:"tool_name"` + Description string `json:"description"` + ConfirmTitle string `json:"confirm_title"` + } + _ = json.Unmarshal(inner.Payload, &req) + what := req.ToolName + if what == "" { + what = req.ConfirmTitle + } + if what == "" { + what = envelope.Type + } + emit(Event{Kind: EventNeedsInput, Tool: what, Text: req.Description}) + + case "agent_event": + var ev struct { + Type string `json:"type"` + Error string `json:"error"` + RunID string `json:"run_id"` + } + if err := json.Unmarshal(inner.Payload, &ev); err == nil && ev.Type == "error" { + if ev.RunID == "" || ev.RunID == runID { + emit(Event{Kind: EventError, Text: ev.Error, IsError: true}) + } + } + + case "run_complete": + var rc struct { + SessionID string `json:"session_id"` + RunID string `json:"run_id"` + Text string `json:"text"` + Error string `json:"error"` + Cancelled bool `json:"cancelled"` + } + if err := json.Unmarshal(inner.Payload, &rc); err != nil { + return false + } + if rc.RunID != runID && !(rc.RunID == "" && rc.SessionID == sessID) { + return false + } + result.Text = rc.Text + result.Err = rc.Error + result.Cancelled = rc.Cancelled + return true + } + return false +} + +func firstLine(s string, n int) string { + s = strings.TrimSpace(s) + if i := strings.IndexByte(s, '\n'); i > 0 { + s = s[:i] + } + if len(s) > n { + s = s[:n-1] + "โ€ฆ" + } + return s +} + +func randomID() string { + buf := make([]byte, 16) + if _, err := rand.Read(buf); err != nil { + return strconv.FormatInt(time.Now().UnixNano(), 36) + } + return hex.EncodeToString(buf) +} diff --git a/internal/harness/crush_serve_test.go b/internal/harness/crush_serve_test.go new file mode 100644 index 0000000..a240f40 --- /dev/null +++ b/internal/harness/crush_serve_test.go @@ -0,0 +1,299 @@ +package harness + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +// fakeCrush is an in-process stand-in for `crush server`, speaking the same +// v1 API surface the adapter uses: health, workspaces, sessions, the SSE +// event stream, agent prompt, cancel and control. +type fakeCrush struct { + listener net.Listener + server *http.Server + + mu sync.Mutex + agentPrompt string + agentRunID string + cancelCalls []string + modelReqs []map[string]any + script func(w io.Writer, r *http.Request, sid, runID string) + agentPosted chan struct{} + postOnce sync.Once +} + +func startFakeCrush(t *testing.T, script func(w io.Writer, r *http.Request, sid, runID string)) *fakeCrush { + t.Helper() + dir := t.TempDir() + sock := filepath.Join(dir, "crush.sock") + ln, err := net.Listen("unix", sock) + if err != nil { + t.Fatal(err) + } + f := &fakeCrush{ + listener: ln, + script: script, + agentPosted: make(chan struct{}), + } + mux := http.NewServeMux() + + writeJSON := func(w http.ResponseWriter, code int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + json.NewEncoder(w).Encode(v) + } + + mux.HandleFunc("GET /v1/health", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + mux.HandleFunc("POST /v1/workspaces", func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + json.NewDecoder(r.Body).Decode(&body) + if body["client_id"] == nil || body["path"] == nil { + writeJSON(w, http.StatusBadRequest, map[string]string{"message": "invalid client_id"}) + return + } + writeJSON(w, http.StatusOK, map[string]any{"id": "ws-1", "path": body["path"]}) + }) + + mux.HandleFunc("POST /v1/workspaces/{id}/sessions", func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, map[string]any{"id": "sess-1"}) + }) + + mux.HandleFunc("POST /v1/workspaces/{id}/current-session", func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, map[string]string{}) + }) + + mux.HandleFunc("POST /v1/workspaces/{id}/config/model", func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + json.NewDecoder(r.Body).Decode(&body) + f.mu.Lock() + f.modelReqs = append(f.modelReqs, body) + f.mu.Unlock() + w.WriteHeader(http.StatusOK) + }) + + mux.HandleFunc("POST /v1/workspaces/{id}/agent/update", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + mux.HandleFunc("GET /v1/workspaces/{id}/events", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("client_id") == "" { + w.WriteHeader(http.StatusBadRequest) + return + } + // Flush headers immediately, exactly like the real server: the + // client's HTTP Do must return before it posts the prompt. + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + select { + case <-f.agentPosted: + case <-r.Context().Done(): + return + } + f.mu.Lock() + runID := f.agentRunID + f.mu.Unlock() + f.script(w, r, "sess-1", runID) + }) + + mux.HandleFunc("POST /v1/workspaces/{id}/agent", func(w http.ResponseWriter, r *http.Request) { + var body struct { + SessionID string `json:"session_id"` + RunID string `json:"run_id"` + Prompt string `json:"prompt"` + } + json.NewDecoder(r.Body).Decode(&body) + f.mu.Lock() + f.agentPrompt, f.agentRunID = body.Prompt, body.RunID + f.mu.Unlock() + f.postOnce.Do(func() { close(f.agentPosted) }) + w.WriteHeader(http.StatusAccepted) + }) + + mux.HandleFunc("POST /v1/workspaces/{id}/agent/sessions/{sid}/cancel", func(w http.ResponseWriter, r *http.Request) { + f.mu.Lock() + f.cancelCalls = append(f.cancelCalls, r.PathValue("sid")) + f.mu.Unlock() + w.WriteHeader(http.StatusOK) + }) + + mux.HandleFunc("POST /v1/control", func(w http.ResponseWriter, r *http.Request) { + io.Copy(io.Discard, r.Body) + w.WriteHeader(http.StatusOK) + }) + + f.server = &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second} + go f.server.Serve(ln) + t.Cleanup(func() { f.server.Close() }) + t.Setenv("TRAUST_CI_CRUSH_SERVER_URL_UNUSED", "") // keep env hermetic intent + t.Setenv("TRAUST_CLI_CRUSH_SERVER_URL", "unix://"+sock) + return f +} + +func sse(w io.Writer, format string, args ...any) { + fmt.Fprintf(w, "data: %s\n\n", fmt.Sprintf(format, args...)) + w.(http.Flusher).Flush() +} + +// happyScript streams the shapes a real crush turn produces. +func happyScript(w io.Writer, r *http.Request, sid, runID string) { + sse(w, `{"type":"message","payload":{"type":"updated","payload":{"id":"m1","role":"assistant","parts":[%s]}}}`, + `{"type":"reasoning","data":{"thinking":"pondering the target","finished_at":1792525268}}`) + sse(w, `{"type":"message","payload":{"type":"updated","payload":{"id":"m1","role":"assistant","parts":[%s,%s]}}}`, + `{"type":"reasoning","data":{"thinking":"pondering the target","finished_at":1792525268}}`, + `{"type":"tool_call","data":{"name":"bash","input":"{\"command\":\"git clone\"}"}}`) + sse(w, `{"type":"session","payload":{"type":"updated","payload":{"id":"%s","prompt_tokens":100,"completion_tokens":50,"cost":0.42}}}`, sid) + sse(w, `{"type":"message","payload":{"type":"updated","payload":{"id":"m1","role":"assistant","parts":[%s,%s,%s]}}}`, + `{"type":"reasoning","data":{"thinking":"pondering the target","finished_at":1792525268}}`, + `{"type":"tool_call","data":{"name":"bash","input":"{\"command\":\"git clone\"}"}}`, + `{"type":"text","data":{"text":"audit complete"}}`) + sse(w, `{"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"%s","run_id":"%s","message_id":"m1","text":"audit complete"}}}`, sid, runID) +} + +func TestCrushServeSessionRun(t *testing.T) { + f := startFakeCrush(t, happyScript) + + var events []Event + res, err := (&Crush{}).RunSession(context.Background(), &RunSpec{ + Prompt: "do the audit", Cwd: t.TempDir(), Env: os.Environ(), + }, func(ev Event) { events = append(events, ev) }) + if err != nil { + t.Fatal(err) + } + if res.Text != "audit complete" || res.Err != "" || res.Cancelled { + t.Fatalf("result = %+v", res) + } + + var sawText, sawThinking, sawTool, sawUsage bool + for _, ev := range events { + switch ev.Kind { + case EventText: + sawText = ev.Text == "audit complete" + case EventThinking: + sawThinking = ev.Text == "pondering the target" + case EventToolCall: + sawTool = ev.Tool == "bash" && strings.Contains(ev.Text, "git clone") + case EventUsage: + sawUsage = ev.TokensIn == 100 && ev.TokensOut == 50 && ev.CostUSD == 0.42 + } + } + if !sawText || !sawThinking || !sawTool || !sawUsage { + t.Errorf("events incomplete: text=%v thinking=%v tool=%v usage=%v (%d events)", sawText, sawThinking, sawTool, sawUsage, len(events)) + } + if len(f.cancelCalls) != 0 { + t.Errorf("unexpected cancel calls: %v", f.cancelCalls) + } +} + +func TestCrushServePermissionFailsClosed(t *testing.T) { + f := startFakeCrush(t, func(w io.Writer, r *http.Request, sid, runID string) { + sse(w, `{"type":"permission_request","payload":{"type":"updated","payload":{"session_id":"%s","tool_name":"write","description":"/etc/hosts"}}}`, sid) + // Hold the stream open until the client disconnects. + <-r.Context().Done() + }) + + var events []Event + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + res, err := (&Crush{}).RunSession(ctx, &RunSpec{Prompt: "p", Cwd: t.TempDir(), Env: os.Environ()}, + func(ev Event) { events = append(events, ev) }) + if err != nil { + t.Fatal(err) + } + needsInput := false + for _, ev := range events { + if ev.Kind == EventNeedsInput && ev.Tool == "write" && ev.Text == "/etc/hosts" { + needsInput = true + } + } + if !needsInput { + t.Errorf("permission request should surface as needs-input: %+v", events) + } + if !res.Cancelled { + t.Errorf("cancelled run should report Cancelled: %+v", res) + } + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + f.mu.Lock() + calls := len(f.cancelCalls) + f.mu.Unlock() + if calls > 0 { + return + } + time.Sleep(50 * time.Millisecond) + } + t.Error("the remote cancel endpoint was never called") +} + +func TestCrushServeRunCompleteError(t *testing.T) { + startFakeCrush(t, func(w io.Writer, r *http.Request, sid, runID string) { + sse(w, `{"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"%s","run_id":"%s","error":"provider 401"}}}`, sid, runID) + }) + res, err := (&Crush{}).RunSession(context.Background(), &RunSpec{Prompt: "p", Cwd: t.TempDir(), Env: os.Environ()}, func(Event) {}) + if err != nil { + t.Fatal(err) + } + if res.Err != "provider 401" { + t.Fatalf("result = %+v", res) + } +} + +func TestCrushServeIgnoresForeignRunComplete(t *testing.T) { + startFakeCrush(t, func(w io.Writer, r *http.Request, sid, runID string) { + // A run_complete for a different run must not terminate ours. + sse(w, `{"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"other","run_id":"not-mine","text":"nope"}}}`) + sse(w, `{"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"%s","run_id":"%s","text":"mine"}}}`, sid, runID) + }) + res, err := (&Crush{}).RunSession(context.Background(), &RunSpec{Prompt: "p", Cwd: t.TempDir(), Env: os.Environ()}, func(Event) {}) + if err != nil { + t.Fatal(err) + } + if res.Text != "mine" { + t.Fatalf("result = %+v", res) + } +} + +func TestCrushServeModelOverride(t *testing.T) { + f := startFakeCrush(t, happyScript) + // Fake `crush models` listing; the runner resolves the floor model here. + dir := t.TempDir() + bin := filepath.Join(dir, "crush-models") + script := "#!/bin/sh\nprintf 'PriceTag/floor-model-x\\nProvider2/other-model\\n'\n" + if err := os.WriteFile(bin, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("TRAUST_CLI_CRUSH_BIN", bin) + + _, err := (&Crush{}).RunSession(context.Background(), &RunSpec{ + Prompt: "p", Cwd: t.TempDir(), Env: os.Environ(), Model: "floor-model-x", + }, func(Event) {}) + if err != nil { + t.Fatal(err) + } + f.mu.Lock() + defer f.mu.Unlock() + if len(f.modelReqs) != 1 { + t.Fatalf("model requests = %+v", f.modelReqs) + } + req := f.modelReqs[0] + if req["model_type"] != "large" || req["scope"] != float64(1) { + t.Errorf("model request = %+v", req) + } + model := req["model"].(map[string]any) + if model["provider"] != "PriceTag" || model["model"] != "floor-model-x" { + t.Errorf("resolved model = %+v", model) + } +} diff --git a/internal/harness/harness.go b/internal/harness/harness.go index 7c4d022..76942c7 100644 --- a/internal/harness/harness.go +++ b/internal/harness/harness.go @@ -22,14 +22,16 @@ const ( EventText EventKind = "text" // EventToolCall is a tool invocation by the agent. EventToolCall EventKind = "tool_call" + // EventThinking is a completed reasoning/thinking block. + EventThinking EventKind = "thinking" // EventUsage is a cumulative token-usage update. EventUsage EventKind = "usage" // EventResult is the final result of the run. EventResult EventKind = "result" // EventError is a stream-reported error. EventError EventKind = "error" - // EventNeedsInput means the agent asked the user a question โ€” impossible - // to satisfy headlessly (exit code 2). + // EventNeedsInput means the agent asked the user a question or raised a + // permission request โ€” impossible to satisfy headlessly (exit code 2). EventNeedsInput EventKind = "needs_input" ) @@ -81,6 +83,32 @@ type RunSpec struct { OutDir string } +// ErrServerUnsupported is returned by a SessionRunner whose managed backend +// cannot start (missing subcommand, old version). The runner falls back to +// the one-shot process path when it sees this. +var ErrServerUnsupported = errors.New("harness server mode unsupported") + +// SessionResult is the terminal outcome of a session-based run. Err and +// Cancelled describe the RUN's fate, not the transport: transport failures +// go into the error return of RunSession. +type SessionResult struct { + Text string + Err string + Cancelled bool +} + +// SessionRunner is implemented by harnesses that execute through a managed +// session backend (crush's `crush server`: workspace, SSE event stream, +// cancel endpoint) instead of a one-shot process. The runner feeds every +// normalized event to emit; RunSession must honor context cancellation by +// cancelling the remote run and returning promptly. +type SessionRunner interface { + // RunSession executes spec.Prompt and emits normalized events until the + // run terminates. Returns ErrServerUnsupported when the backend cannot + // start, so the runner can fall back to the process path. + RunSession(ctx context.Context, spec *RunSpec, emit func(Event)) (SessionResult, error) +} + // Harness is one agent harness adapter. type Harness interface { ID() string diff --git a/internal/runner/runner.go b/internal/runner/runner.go index 237ddae..808bb59 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -189,52 +189,23 @@ func Run(ctx context.Context, opts Options, sink Sink) (*RunResult, error) { // Launch. runCtx, cancel := context.WithCancel(ctx) defer cancel() - cmd := h.Command(runCtx, &harness.RunSpec{ + spec := &harness.RunSpec{ Prompt: prompt, Model: pf.Model, Cwd: root, Env: state.ChildEnv(), AllowedTools: skill.AllowedTools, MaxTurns: opts.MaxTurns, UnsafeFullPermission: opts.UnsafeFullPermission, OutDir: journal.Dir, - }) - stderrPath := filepath.Join(journal.Dir, "harness-stderr.log") - stderrFile, err := os.Create(stderrPath) - if err != nil { - return nil, fmt.Errorf("create stderr log: %w", err) - } - defer stderrFile.Close() - cmd.Stderr = stderrFile - stdout, err := cmd.StdoutPipe() - if err != nil { - return nil, err - } - if err := cmd.Start(); err != nil { - return nil, fmt.Errorf("start %s: %w", h.Binary(), err) } - // Stream loop. - streaming := h.Capabilities().StreamJSON - var plainOut strings.Builder - scanner := bufio.NewScanner(stdout) - scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + // handleEvent folds one normalized event into journal, sink, result and + // the budget bands โ€” shared by the session and process paths. band := BandWithin - for scanner.Scan() { - line := scanner.Text() - journal.Raw(line) - if !streaming { - plainOut.WriteString(line) - plainOut.WriteString("\n") - sink.Raw(line) - } else if ev, ok := h.ParseEvent(scanner.Bytes()); ok { - journal.Event(string(ev.Kind), ev.Text, ev.Tool, ev.TokensIn, ev.TokensOut, ev.CostUSD, ev.IsError) - sink.Event(ev) - applyEvent(res, ev) - if ev.Kind == harness.EventNeedsInput { - res.NeedsInput = true - cancel() - } - } else { - sink.Raw(line) + handleEvent := func(ev harness.Event) { + journal.Event(string(ev.Kind), ev.Text, ev.Tool, ev.TokensIn, ev.TokensOut, ev.CostUSD, ev.IsError) + sink.Event(ev) + applyEvent(res, ev) + if ev.Kind == harness.EventNeedsInput { + res.NeedsInput = true + cancel() } - - // Budget bands on cumulative output tokens. if policy != nil && res.Usage.TokensOut > 0 { b := policy.BandFor(skill.Name, res.Usage.TokensOut) if b != band { @@ -248,10 +219,77 @@ func Run(ctx context.Context, opts Options, sink Sink) (*RunResult, error) { } } } - scanErr := scanner.Err() - waitErr := cmd.Wait() - if exitErr, ok := waitErr.(*exec.ExitError); ok { - res.HarnessExit = exitErr.ExitCode() + + // Session mode first: harnesses with a managed backend (crush server) + // stream structured events, usage and a clean cancel. Falls back to the + // one-shot process path when the backend cannot start. + var sessErr error + sessionDone := false + if sr, ok := h.(harness.SessionRunner); ok { + sres, err := sr.RunSession(runCtx, spec, handleEvent) + switch { + case errors.Is(err, harness.ErrServerUnsupported): + sink.Progress("note: crush server mode unavailable โ€” falling back to `crush run` (plain text, no live events)") + case err != nil: + sessErr = err + sessionDone = true + default: + sessionDone = true + if sres.Text != "" { + res.ResultText = sres.Text + } + if sres.Err != "" { + res.HarnessError = sres.Err + } + if sres.Cancelled { + res.HarnessExit = -1 + } + } + } + + var waitErr, scanErr error + stderrPath := filepath.Join(journal.Dir, "harness-stderr.log") + var plainOut strings.Builder + if !sessionDone && sessErr == nil { + cmd := h.Command(runCtx, spec) + stderrFile, err := os.Create(stderrPath) + if err != nil { + return nil, fmt.Errorf("create stderr log: %w", err) + } + cmd.Stderr = stderrFile + stdout, err := cmd.StdoutPipe() + if err != nil { + stderrFile.Close() + return nil, err + } + if err := cmd.Start(); err != nil { + stderrFile.Close() + return nil, fmt.Errorf("start %s: %w", h.Binary(), err) + } + + // Stream loop. + streaming := h.Capabilities().StreamJSON + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + for scanner.Scan() { + line := scanner.Text() + journal.Raw(line) + if !streaming { + plainOut.WriteString(line) + plainOut.WriteString("\n") + sink.Raw(line) + } else if ev, ok := h.ParseEvent(scanner.Bytes()); ok { + handleEvent(ev) + } else { + sink.Raw(line) + } + } + scanErr = scanner.Err() + waitErr = cmd.Wait() + stderrFile.Close() + if exitErr, ok := waitErr.(*exec.ExitError); ok { + res.HarnessExit = exitErr.ExitCode() + } } switch { @@ -260,11 +298,18 @@ func Run(ctx context.Context, opts Options, sink Sink) (*RunResult, error) { sink.Progress(fmt.Sprintf("budget refuse band hit (%d output tokens, cap %d) โ€” harness terminated; re-run with --budget-override \"\" if this is expected", res.Usage.TokensOut, res.BudgetCap)) case res.NeedsInput: res.ExitCode = ExitNeedsInput - sink.Progress("the skill asked for interactive input, which a headless run cannot provide โ€” run it interactively in your harness and answer its questions, then re-run") + sink.Progress("the skill asked for interactive input or a permission grant, which a headless run cannot provide โ€” run it interactively in your harness and answer its questions, then re-run") case ctx.Err() != nil: res.Aborted = true res.ExitCode = ExitError sink.Progress("run aborted") + case sessErr != nil: + res.ExitCode = ExitError + res.HarnessError = sessErr.Error() + sink.Progress("harness session failed: " + res.HarnessError) + case res.HarnessError != "": + res.ExitCode = ExitError + sink.Progress("harness run failed: " + res.HarnessError) case waitErr != nil: res.ExitCode = ExitError res.HarnessError = waitErr.Error() @@ -277,7 +322,7 @@ func Run(ctx context.Context, opts Options, sink Sink) (*RunResult, error) { sink.Progress(fmt.Sprintf("stream read error: %v", scanErr)) } - if !streaming && res.ResultText == "" { + if !sessionDone && !h.Capabilities().StreamJSON && res.ResultText == "" { res.ResultText = strings.TrimRight(plainOut.String(), "\n") } diff --git a/internal/runner/session_test.go b/internal/runner/session_test.go new file mode 100644 index 0000000..967c3ed --- /dev/null +++ b/internal/runner/session_test.go @@ -0,0 +1,180 @@ +package runner + +import ( + "context" + "encoding/json" + "fmt" + "net" + "net/http" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/traust-security/traust-cli/internal/catalog" + "github.com/traust-security/traust-cli/internal/harness" +) + +// fakeServeCrush runs a minimal crush-server v1 API on a Unix socket and +// points the crush adapter at it via TRAUST_CLI_CRUSH_SERVER_URL. +func fakeServeCrush(t *testing.T, script string) { + t.Helper() + mux := http.NewServeMux() + var posted atomic.Bool + var runID atomic.Value + + mux.HandleFunc("GET /v1/health", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) }) + mux.HandleFunc("POST /v1/workspaces", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"id": "ws-1"}) + }) + mux.HandleFunc("POST /v1/workspaces/{id}/sessions", func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]string{"id": "sess-1"}) + }) + mux.HandleFunc("POST /v1/workspaces/{id}/current-session", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + }) + mux.HandleFunc("POST /v1/workspaces/{id}/config/model", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + }) + mux.HandleFunc("POST /v1/workspaces/{id}/agent/update", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + }) + mux.HandleFunc("POST /v1/workspaces/{id}/agent", func(w http.ResponseWriter, r *http.Request) { + var body struct { + RunID string `json:"run_id"` + } + json.NewDecoder(r.Body).Decode(&body) + runID.Store(body.RunID) + posted.Store(true) + w.WriteHeader(202) + }) + mux.HandleFunc("POST /v1/workspaces/{id}/agent/sessions/{sid}/cancel", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + }) + mux.HandleFunc("GET /v1/workspaces/{id}/events", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(200) + w.(http.Flusher).Flush() + for !posted.Load() { + select { + case <-r.Context().Done(): + return + case <-time.After(20 * time.Millisecond): + } + } + fmt.Fprintf(w, "data: %s\n\n", strings.ReplaceAll(script, "{RUN_ID}", runID.Load().(string))) + w.(http.Flusher).Flush() + <-r.Context().Done() + }) + + ln, err := net.Listen("unix", filepath.Join(t.TempDir(), "s.sock")) + if err != nil { + t.Fatal(err) + } + srv := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second} + go srv.Serve(ln) + t.Cleanup(func() { srv.Close() }) + t.Setenv("TRAUST_CLI_CRUSH_SERVER_URL", "unix://"+ln.Addr().String()) +} + +const happyServeScript = `{"type":"message","payload":{"type":"updated","payload":{"id":"m1","role":"assistant","parts":[{"type":"tool_call","data":{"name":"bash","input":"{\"command\":\"ls\"}"}}]}}} +data: {"type":"session","payload":{"type":"updated","payload":{"id":"sess-1","prompt_tokens":120,"completion_tokens":60,"cost":1.5}}} +data: {"type":"message","payload":{"type":"updated","payload":{"id":"m1","role":"assistant","parts":[{"type":"tool_call","data":{"name":"bash","input":"{\"command\":\"ls\"}"}},{"type":"text","data":{"text":"audit complete"}}]}}} +data: {"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"sess-1","run_id":"{RUN_ID}","text":"audit complete"}}}` + +func TestRunSessionModeHappyPath(t *testing.T) { + fakeServeCrush(t, happyServeScript) + // A fake `crush models` so the registry floor model resolves and the + // model-check pre-flight passes. + bin := filepath.Join(t.TempDir(), "crush") + if err := os.WriteFile(bin, []byte("#!/bin/sh\nprintf 'PriceTag/floor-model-x\\n'\n"), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("TRAUST_CLI_CRUSH_BIN", bin) + + state := testState(t) + client := &fakeAPI{} + sink := &recSink{} + + // A registry role maps the skill to the floor model. + skill := testSkill() + skill.Name = "vuln-scan" + skill.Standalone = &catalog.StandaloneInfo{Role: "vuln-scan", MinModelClass: "opus"} + + res, err := Run(context.Background(), Options{ + State: state, API: client, Harness: &harness.Crush{}, + Skill: skill, Args: []string{"https://github.com/org/repo"}, + DeclareSpend: true, + }, sink) + if err != nil { + t.Fatal(err) + } + if res.ExitCode != ExitOK { + t.Fatalf("exit = %d (%+v)", res.ExitCode, res) + } + if res.ResultText != "audit complete" { + t.Errorf("result text = %q", res.ResultText) + } + if res.Usage.TokensIn != 120 || res.Usage.TokensOut != 60 || res.Usage.CostUSD != 1.5 { + t.Errorf("usage = %+v", res.Usage) + } + if len(client.spendRows) != 1 || client.spendRows[0].TokensOut != 60 { + t.Errorf("spend rows = %+v (session usage should be declared)", client.spendRows) + } + var sawTool, sawText bool + for _, ev := range sink.events { + if ev.Kind == harness.EventToolCall && ev.Tool == "bash" { + sawTool = true + } + if ev.Kind == harness.EventText { + sawText = true + } + } + if !sawTool || !sawText { + t.Errorf("live events missing: tool=%v text=%v", sawTool, sawText) + } + if res.HarnessExit != 0 { + t.Errorf("harness exit = %d", res.HarnessExit) + } +} + +func TestRunSessionModeNeedsInput(t *testing.T) { + fakeServeCrush(t, `{"type":"permission_request","payload":{"type":"updated","payload":{"session_id":"sess-1","tool_name":"write","description":"/etc/shadow"}}} +data: {"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"sess-1","run_id":"{RUN_ID}","cancelled":true}}}`) + state := testState(t) + + skill := testSkill() + skill.Standalone = &catalog.StandaloneInfo{MinModelClass: "opus"} + + res, err := Run(context.Background(), Options{ + State: state, API: &fakeAPI{}, Harness: &harness.Crush{}, + Skill: skill, DeclareSpend: true, + }, nil) + if err != nil { + t.Fatal(err) + } + if res.ExitCode != ExitNeedsInput || !res.NeedsInput { + t.Fatalf("exit = %d, needsInput = %v", res.ExitCode, res.NeedsInput) + } +} + +func TestRunSessionModeRunError(t *testing.T) { + fakeServeCrush(t, `{"type":"run_complete","payload":{"type":"updated","payload":{"session_id":"sess-1","run_id":"{RUN_ID}","error":"provider 401"}}}`) + state := testState(t) + + skill := testSkill() + skill.Standalone = &catalog.StandaloneInfo{MinModelClass: "opus"} + + res, err := Run(context.Background(), Options{ + State: state, API: &fakeAPI{}, Harness: &harness.Crush{}, + Skill: skill, DeclareSpend: true, + }, nil) + if err != nil { + t.Fatal(err) + } + if res.ExitCode != ExitError || res.HarnessError != "provider 401" { + t.Fatalf("exit = %d, harnessError = %q", res.ExitCode, res.HarnessError) + } +} diff --git a/internal/tui/monitor.go b/internal/tui/monitor.go index 4963fea..1e7bdb8 100644 --- a/internal/tui/monitor.go +++ b/internal/tui/monitor.go @@ -207,8 +207,14 @@ func (m *monitorModel) applyEvent(ev harness.Event) { switch ev.Kind { case harness.EventText: m.pushLine(ev.Text) + case harness.EventThinking: + m.pushLine(m.styles.Muted.Italic(true).Render("ยท thinking: " + ev.Text)) case harness.EventToolCall: - m.pushLine(m.styles.Accent.Render("โ†’ " + ev.Tool)) + line := "โ†’ " + ev.Tool + if ev.Text != "" { + line += " " + ev.Text + } + m.pushLine(m.styles.Accent.Render(line)) case harness.EventUsage: m.tokensIn, m.tokensOut = ev.TokensIn, ev.TokensOut case harness.EventResult: @@ -222,7 +228,11 @@ func (m *monitorModel) applyEvent(ev harness.Event) { m.pushLine(m.styles.Err.Render("โœ— " + ev.Text)) case harness.EventNeedsInput: m.failed = true - m.pushLine(m.styles.Err.Render("โœ— the skill asked for interactive input โ€” run it interactively in your harness")) + what := ev.Tool + if what == "" { + what = ev.Text + } + m.pushLine(m.styles.Err.Render("โœ— needs input that a headless run cannot provide: " + what)) } } @@ -231,9 +241,7 @@ func (m *monitorModel) applyEvent(ev harness.Event) { // complete at any terminal size. func (m *monitorModel) pushLine(line string) { line = strings.ReplaceAll(line, "\r", "") - for _, l := range strings.Split(line, "\n") { - m.lines = append(m.lines, l) - } + m.lines = append(m.lines, strings.Split(line, "\n")...) if len(m.lines) > 1000 { m.lines = m.lines[len(m.lines)-1000:] }