90b1a2a · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat(server): make server prompts independent of client connections

90b1a2aChristian Rocha5 months ago

internal/agent/agent.go+22−7

compare with main

unifiedsplit

@@ -305,7 +305,13 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

305 // (the pubsub broker fan-in does not serialize publishes from

306 // different upstream brokers).

307 defer func() {

308+ // Use a context detached from the run context: workspace

309+ // shutdown cancels ctx before this goroutine returns, but the

310+ // buffered streaming deltas must still land before the DB is

311+ // closed. A short timeout bounds the flush.

312+ flushCtx, flushCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)

313+ defer flushCancel()

308− if flushErr := a.messages.FlushAll(ctx); flushErr != nil {

314+ if flushErr := a.messages.FlushAll(flushCtx); flushErr != nil {

309 slog.Error("Failed to flush pending message updates after run", "error", flushErr)

315 slog.Error("Failed to flush pending message updates after run", "error", flushErr)

310 }

316 }

311 if skipRunComplete {

317 if skipRunComplete {

@@ -577,11 +583,20 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

577 if currentAssistant == nil {

583 if currentAssistant == nil {

578 return result, err

584 return result, err

579 }

585 }

586+ // Persist final state with a context detached from the run

587+ // context. The run context (ctx) is derived from the

588+ // workspace context, which workspace shutdown cancels before

589+ // agent goroutines finish; using ctx here would drop the

590+ // final assistant state. WithoutCancel keeps the values

591+ // (e.g. session ID) while ignoring cancellation, and a short

592+ // timeout bounds the cleanup writes.

593+ cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)

594+ defer cleanupCancel()

580 // Ensure we finish thinking on error to close the reasoning state.

595 // Ensure we finish thinking on error to close the reasoning state.

581 currentAssistant.FinishThinking()

596 currentAssistant.FinishThinking()

582 toolCalls := currentAssistant.ToolCalls()

597 toolCalls := currentAssistant.ToolCalls()

583− // INFO: we use the parent context here because the genCtx has been cancelled.

598+ // INFO: we use the cleanup context here because the genCtx has been cancelled.

584− msgs, createErr := a.messages.List(ctx, currentAssistant.SessionID)

599+ msgs, createErr := a.messages.List(cleanupCtx, currentAssistant.SessionID)

585 if createErr != nil {

600 if createErr != nil {

586 return nil, createErr

601 return nil, createErr

587 }

602 }

@@ -590,7 +605,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

590 tc.Finished = true

605 tc.Finished = true

591 tc.Input = "{}"

606 tc.Input = "{}"

592 currentAssistant.AddToolCall(tc)

607 currentAssistant.AddToolCall(tc)

593− updateErr := a.messages.Update(ctx, *currentAssistant)

608+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)

594 if updateErr != nil {

609 if updateErr != nil {

595 return nil, updateErr

610 return nil, updateErr

596 }

611 }

@@ -623,7 +638,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

623 Content: content,

638 Content: content,

624 IsError: true,

639 IsError: true,

625 }

640 }

626− _, createErr = a.messages.Create(ctx, currentAssistant.SessionID, message.CreateMessageParams{

641+ _, createErr = a.messages.Create(cleanupCtx, currentAssistant.SessionID, message.CreateMessageParams{

627 Role: message.Tool,

642 Role: message.Tool,

628 Parts: []message.ContentPart{

643 Parts: []message.ContentPart{

629 toolResult,

644 toolResult,

@@ -670,9 +685,9 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *

670 } else {

685 } else {

671 currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error())

686 currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error())

672 }

687 }

673− // Note: we use the parent context here because the genCtx has been

688+ // Note: we use the cleanup context here because the genCtx has been

674 // cancelled.

689 // cancelled.

675− updateErr := a.messages.Update(ctx, *currentAssistant)

690+ updateErr := a.messages.Update(cleanupCtx, *currentAssistant)

676 if updateErr != nil {

691 if updateErr != nil {

677 return nil, updateErr

692 return nil, updateErr

678 }

693 }

internal/backend/backend.go+56−1

@@ -34,6 +34,7 @@ var (

34 ErrUnknownCommand = errors.New("unknown command")

35 ErrInvalidClientID = errors.New("invalid client_id")

36 ErrClientNotAttached = errors.New("client not attached")

37+ ErrWorkspaceClosing = errors.New("workspace closing")

37)

38)

39// DefaultCreateGrace is the window in which a client must open an SSE

40// DefaultCreateGrace is the window in which a client must open an SSE

@@ -108,6 +109,23 @@ type Workspace struct {

108 // with fallback to the cleaned absolute path.

109 // with fallback to the cleaned absolute path.

109 resolvedPath string

110 resolvedPath string

110

111

112+ // ctx is the workspace-scoped run context. It is derived from

113+ // the backend context in CreateWorkspace and lives for the

114+ // lifetime of the workspace; cancel tears it down. Agent runs

115+ // dispatched on behalf of this workspace are bound to ctx so

116+ // their lifetime is owned by the workspace, not by any single

117+ // client's HTTP request.

118+ ctx context.Context

119+ cancel context.CancelFunc

120+

121+ // runMu guards closing and gates dispatch of new agent runs.

122+ // closing is set by Shutdown so no new runs are accepted once

123+ // teardown has begun. runWG tracks dispatched agent goroutines

124+ // so Shutdown can wait for them to return before app cleanup.

125+ runMu sync.Mutex

126+ closing bool

127+ runWG sync.WaitGroup

128+

111 // clientsMu guards clients. It is held only briefly (no IO).

129 // clientsMu guards clients. It is held only briefly (no IO).

112 clientsMu sync.Mutex

130 clientsMu sync.Mutex

113 // clients tracks each client's claim on this workspace. Refcount

131 // clients tracks each client's claim on this workspace. Refcount

@@ -122,7 +140,7 @@ type Workspace struct {

122}

140}

123

141

124// invokeShutdown calls the workspace shutdown hook if set, falling

142// invokeShutdown calls the workspace shutdown hook if set, falling

125−// back to the embedded [app.App.Shutdown] when not.

143+// back to the workspace [Workspace.Shutdown] wrapper when not.

126func (w *Workspace) invokeShutdown() {

144func (w *Workspace) invokeShutdown() {

127 if w.shutdownFn != nil {

145 if w.shutdownFn != nil {

128 w.shutdownFn()

146 w.shutdownFn()

@@ -133,6 +151,40 @@ func (w *Workspace) invokeShutdown() {

133 }

151 }

134}

152}

135

153

154+// Shutdown tears the workspace down in an order that is safe for

155+// agent runs whose lifetime is bound to the workspace context. It

156+// shadows the promoted [app.App.Shutdown] so callers reaching

157+// ws.Shutdown() always observe this ordering:

158+//

159+// 1. Mark the workspace closing so no new agent runs are accepted.

160+// 2. Cancel the workspace run context so any dispatched goroutine

161+// that has not yet registered its per-session cancel still

162+// observes cancellation.

163+// 3. Cancel active coordinator work for runs that already

164+// registered their per-session cancel function.

165+// 4. Wait for dispatched agent goroutines to return.

166+// 5. Run the embedded [app.App.Shutdown] cleanup (DB, LSP, etc).

167+//

168+// CancelAll is idempotent, so the second call inside app.App.Shutdown

169+// is harmless; the important guarantee is that cancel -> CancelAll ->

170+// runWG.Wait completes before the embedded cleanup touches the DB.

171+func (w *Workspace) Shutdown() {

172+ w.runMu.Lock()

173+ w.closing = true

174+ w.runMu.Unlock()

175+

176+ if w.cancel != nil {

177+ w.cancel()

178+ }

179+ if w.App != nil && w.AgentCoordinator != nil {

180+ w.AgentCoordinator.CancelAll()

181+ }

182+ w.runWG.Wait()

183+ if w.App != nil {

184+ w.App.Shutdown()

185+ }

186+}

187+

136// New creates a new [Backend].

188// New creates a new [Backend].

137func New(ctx context.Context, cfg *config.ConfigStore, shutdownFn ShutdownFunc) *Backend {

189func New(ctx context.Context, cfg *config.ConfigStore, shutdownFn ShutdownFunc) *Backend {

138 return &Backend{

190 return &Backend{

@@ -247,6 +299,7 @@ func (b *Backend) CreateWorkspace(args proto.Workspace) (*Workspace, proto.Works

247 return nil, proto.Workspace{}, fmt.Errorf("failed to create app workspace: %w", err)

299 return nil, proto.Workspace{}, fmt.Errorf("failed to create app workspace: %w", err)

248 }

300 }

249

301

302+ wsCtx, wsCancel := context.WithCancel(b.ctx)

250 ws := &Workspace{

303 ws := &Workspace{

251 App: appWorkspace,

304 App: appWorkspace,

252 ID: id,

305 ID: id,

@@ -255,6 +308,8 @@ func (b *Backend) CreateWorkspace(args proto.Workspace) (*Workspace, proto.Works

255 Env: args.Env,

308 Env: args.Env,

256 Skills: skillsMgr,

309 Skills: skillsMgr,

257 resolvedPath: key,

310 resolvedPath: key,

311+ ctx: wsCtx,

312+ cancel: wsCancel,

258 clients: make(map[string]*clientState),

313 clients: make(map[string]*clientState),

259 }

314 }

260

315