Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 23 additions & 11 deletions internal/acp/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"strings"
"sync"

Expand Down Expand Up @@ -83,7 +84,7 @@
}

// Serve runs the connection read loop until the stream closes or ctx is done.
func (a *Agent) Serve(ctx context.Context) error { return a.conn.Serve(ctx) }

Check failure on line 87 in internal/acp/agent.go

View workflow job for this annotation

GitHub Actions / Security & code health

unreachable func: Agent.Serve

// ---- initialize ----

Expand Down Expand Up @@ -159,8 +160,11 @@
// Load history BEFORE publishing the session so no concurrent prompt observes
// a half-initialized session (registerSession sets history under the lock and
// reuses an already-live session rather than orphaning its in-flight turn).
history := a.loadHistory(meta.SessionID)
history, historyErr := a.loadHistory(meta.SessionID)
sess := a.registerSession(meta.SessionID, root, history)
if historyErr != nil {
(&notifier{conn: a.conn, sessionID: sess.id}).warning("Warning: ACP session history could not be loaded: " + historyErr.Error())
}
return LoadSessionResult{
Modes: a.modeState(sess),
}, nil
Expand Down Expand Up @@ -253,7 +257,9 @@
if stopErr != nil {
return "", RPCError(codeInternalError, stopErr.Error())
}
a.persistTurn(sess, userText, result.FinalAnswer)
if err := a.persistTurn(sess, userText, result.FinalAnswer); err != nil {
note.warning("Warning: ACP session history was not fully persisted: " + err.Error())
}
return reason, nil
}

Expand Down Expand Up @@ -386,29 +392,35 @@

// ---- persistence + continuity ----

func (a *Agent) persistTurn(sess *acpSession, user, assistant string) {
func (a *Agent) persistTurn(sess *acpSession, user, assistant string) error {
var persistErr error
if a.deps.Store != nil {
_, _ = a.deps.Store.AppendEvent(sess.id, sessions.AppendEventInput{
if _, err := a.deps.Store.AppendEvent(sess.id, sessions.AppendEventInput{
Type: sessions.EventMessage,
Payload: map[string]any{"role": "user", "content": user},
})
}); err != nil {
persistErr = errors.Join(persistErr, fmt.Errorf("user message: %w", err))
}
if assistant != "" {
_, _ = a.deps.Store.AppendEvent(sess.id, sessions.AppendEventInput{
if _, err := a.deps.Store.AppendEvent(sess.id, sessions.AppendEventInput{
Type: sessions.EventMessage,
Payload: map[string]any{"role": "assistant", "content": assistant},
})
}); err != nil {
persistErr = errors.Join(persistErr, fmt.Errorf("assistant message: %w", err))
}
}
}
sess.appendHistory(turnRecord{user: user, assistant: assistant})
return persistErr
}

func (a *Agent) loadHistory(sessionID string) []turnRecord {
func (a *Agent) loadHistory(sessionID string) ([]turnRecord, error) {
if a.deps.Store == nil {
return nil
return nil, nil
}
events, err := a.deps.Store.ReadEvents(sessionID)
if err != nil {
return nil
return nil, err
}
var records []turnRecord
var pendingUser string
Expand Down Expand Up @@ -444,7 +456,7 @@
if havePending {
records = append(records, turnRecord{user: pendingUser})
}
return records
return records, nil
}

// buildPrompt prepends prior conversation as context, since agent.Run drives a
Expand Down
89 changes: 83 additions & 6 deletions internal/acp/agent_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import (
"encoding/json"
"fmt"
"io"
"os"
"path/filepath"
"strings"
"testing"
"time"
Expand Down Expand Up @@ -63,9 +65,10 @@ func testDeps(t *testing.T) Deps {
// clientHarness wires a client Conn to an Agent over in-memory pipes and collects
// session/update text chunks.
type clientHarness struct {
client *Conn
updates chan string
stop func()
client *Conn
updates chan string
thoughts chan string
stop func()
}

func newHarness(t *testing.T, deps Deps) *clientHarness {
Expand All @@ -76,7 +79,7 @@ func newHarness(t *testing.T, deps Deps) *clientHarness {
client := NewConn(br, bw)
a := NewAgent(agentConn, deps)

h := &clientHarness{client: client, updates: make(chan string, 128)}
h := &clientHarness{client: client, updates: make(chan string, 128), thoughts: make(chan string, 128)}
client.HandleNotify(MethodSessionUpdate, func(_ context.Context, params json.RawMessage) {
var probe struct {
Update struct {
Expand All @@ -89,8 +92,11 @@ func newHarness(t *testing.T, deps Deps) *clientHarness {
if json.Unmarshal(params, &probe) != nil {
return
}
if probe.Update.SessionUpdate == UpdateAgentMessageChunk {
switch probe.Update.SessionUpdate {
case UpdateAgentMessageChunk:
h.updates <- probe.Update.Content.Text
case UpdateAgentThoughtChunk:
h.thoughts <- probe.Update.Content.Text
}
})

Expand Down Expand Up @@ -154,6 +160,72 @@ func TestACPEndToEndPrompt(t *testing.T) {
}
}

func TestACPPromptWarnsWhenPersistenceFails(t *testing.T) {
deps := testDeps(t)
h := newHarness(t, deps)
defer h.stop()

ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()

var newRes NewSessionResult
if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil {
t.Fatalf("session/new: %v", err)
}
eventsPath := filepath.Join(deps.Store.RootDir, newRes.SessionID, sessions.EventsFile)
if err := os.Remove(eventsPath); err != nil {
t.Fatalf("remove events file: %v", err)
}
if err := os.Mkdir(eventsPath, 0o700); err != nil {
t.Fatalf("replace events file with directory: %v", err)
}

var promptRes PromptResult
if err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{
SessionID: newRes.SessionID,
Prompt: []ContentBlock{TextBlock("hi")},
}, &promptRes); err != nil {
t.Fatalf("session/prompt should continue after persistence failure: %v", err)
}
if promptRes.StopReason != StopEndTurn {
t.Fatalf("stopReason = %q, want %q", promptRes.StopReason, StopEndTurn)
}
if got := drainContains(t, h.updates, "Hello from ZERO"); got == "" {
t.Fatal("expected assistant response to stream despite persistence failure")
}
if got := drainContains(t, h.thoughts, "ACP session history was not fully persisted"); got == "" {
t.Fatal("expected persistence warning thought chunk")
}
}

func TestACPLoadWarnsWhenHistoryReadFails(t *testing.T) {
deps := testDeps(t)
meta, err := deps.Store.Create(sessions.CreateInput{Title: "corrupt", Cwd: t.TempDir()})
if err != nil {
t.Fatalf("create session: %v", err)
}
eventsPath := filepath.Join(deps.Store.RootDir, meta.SessionID, sessions.EventsFile)
if err := os.WriteFile(eventsPath, []byte("{not-json}\n"), 0o600); err != nil {
t.Fatalf("corrupt events file: %v", err)
}

h := newHarness(t, deps)
defer h.stop()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

var loadRes LoadSessionResult
if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: meta.SessionID}, &loadRes); err != nil {
t.Fatalf("session/load should continue after history read failure: %v", err)
}
if loadRes.Modes == nil || loadRes.Modes.CurrentModeID != string(agent.PermissionModeAuto) {
t.Fatalf("expected auto mode, got %+v", loadRes.Modes)
}
if got := drainContains(t, h.thoughts, "ACP session history could not be loaded"); got == "" {
t.Fatal("expected load warning thought chunk")
}
}

func TestACPUnknownSessionPromptErrors(t *testing.T) {
h := newHarness(t, testDeps(t))
defer h.stop()
Expand Down Expand Up @@ -244,14 +316,19 @@ func TestACPRejectsInvalidCwd(t *testing.T) {

// drainText collects streamed chunks for a short window and concatenates them.
func drainText(t *testing.T, ch <-chan string) string {
t.Helper()
return drainContains(t, ch, "Hello from ZERO")
}

func drainContains(t *testing.T, ch <-chan string, needle string) string {
t.Helper()
var b strings.Builder
deadline := time.After(2 * time.Second)
for {
select {
case s := <-ch:
b.WriteString(s)
if strings.Contains(b.String(), "Hello from ZERO") {
if strings.Contains(b.String(), needle) {
return b.String()
}
case <-deadline:
Expand Down
4 changes: 4 additions & 0 deletions internal/acp/translate.go
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,10 @@ func (n *notifier) thought(delta string) {
}
}

func (n *notifier) warning(message string) {
n.thought(message)
}

func (n *notifier) toolCall(call agent.ToolCall) { n.send(toolCallStart(call)) }
func (n *notifier) toolResult(result agent.ToolResult) { n.send(toolCallResult(result)) }

Expand Down
Loading