From f2920739d8d229bede69153c31784e8a9dada3ff Mon Sep 17 00:00:00 2001 From: chenhg5 Date: Wed, 4 Mar 2026 22:02:50 +0800 Subject: [PATCH] feat: add timeout handling for agent session closure Added timeout mechanisms with logging for graceful session closure across multiple agent types (claudecode, codex, cursor, gemini, opencode, qoder). Modified engine cleanup to handle timeouts and improve logging. Sessions now timeout after 8-10 seconds with appropriate warnings/errors. generated by llmgit Co-Authored-By: Claude --- agent/claudecode/session.go | 14 +++++++++-- agent/codex/session.go | 12 +++++++++- agent/cursor/session.go | 12 +++++++++- agent/gemini/session.go | 12 +++++++++- agent/opencode/session.go | 12 +++++++++- agent/qoder/session.go | 12 +++++++++- core/engine.go | 47 +++++++++++++++++++++++++++---------- 7 files changed, 102 insertions(+), 19 deletions(-) diff --git a/agent/claudecode/session.go b/agent/claudecode/session.go index 8badc68..d8d4f3c 100644 --- a/agent/claudecode/session.go +++ b/agent/claudecode/session.go @@ -466,8 +466,18 @@ func (cs *claudeSession) Alive() bool { func (cs *claudeSession) Close() error { cs.cancel() - <-cs.done - return nil + + select { + case <-cs.done: + return nil + case <-time.After(8 * time.Second): + slog.Warn("claudeSession: graceful close timed out, killing process") + if cs.cmd != nil && cs.cmd.Process != nil { + _ = cs.cmd.Process.Kill() + } + <-cs.done + return nil + } } // filterEnv returns a copy of env with entries matching the given key removed. diff --git a/agent/codex/session.go b/agent/codex/session.go index 31ec3ac..0bae967 100644 --- a/agent/codex/session.go +++ b/agent/codex/session.go @@ -13,6 +13,7 @@ import ( "strings" "sync" "sync/atomic" + "time" "unicode/utf8" "github.com/chenhg5/cc-connect/core" @@ -291,7 +292,16 @@ func (cs *codexSession) Alive() bool { func (cs *codexSession) Close() error { cs.alive.Store(false) cs.cancel() - cs.wg.Wait() + done := make(chan struct{}) + go func() { + cs.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(8 * time.Second): + slog.Warn("codexSession: close timed out, abandoning wg.Wait") + } close(cs.events) return nil } diff --git a/agent/cursor/session.go b/agent/cursor/session.go index e844e8b..9ae13d3 100644 --- a/agent/cursor/session.go +++ b/agent/cursor/session.go @@ -13,6 +13,7 @@ import ( "strings" "sync" "sync/atomic" + "time" "unicode/utf8" "github.com/chenhg5/cc-connect/core" @@ -382,7 +383,16 @@ func (cs *cursorSession) Alive() bool { func (cs *cursorSession) Close() error { cs.alive.Store(false) cs.cancel() - cs.wg.Wait() + done := make(chan struct{}) + go func() { + cs.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(8 * time.Second): + slog.Warn("cursorSession: close timed out, abandoning wg.Wait") + } close(cs.events) return nil } diff --git a/agent/gemini/session.go b/agent/gemini/session.go index 086b8e9..498a43d 100644 --- a/agent/gemini/session.go +++ b/agent/gemini/session.go @@ -13,6 +13,7 @@ import ( "strings" "sync" "sync/atomic" + "time" "unicode/utf8" "github.com/chenhg5/cc-connect/core" @@ -370,7 +371,16 @@ func (gs *geminiSession) Alive() bool { func (gs *geminiSession) Close() error { gs.alive.Store(false) gs.cancel() - gs.wg.Wait() + done := make(chan struct{}) + go func() { + gs.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(8 * time.Second): + slog.Warn("geminiSession: close timed out, abandoning wg.Wait") + } close(gs.events) return nil } diff --git a/agent/opencode/session.go b/agent/opencode/session.go index 16be023..31d354e 100644 --- a/agent/opencode/session.go +++ b/agent/opencode/session.go @@ -12,6 +12,7 @@ import ( "os/exec" "sync" "sync/atomic" + "time" "unicode/utf8" "github.com/chenhg5/cc-connect/core" @@ -305,7 +306,16 @@ func (s *opencodeSession) Alive() bool { func (s *opencodeSession) Close() error { s.alive.Store(false) s.cancel() - s.wg.Wait() + done := make(chan struct{}) + go func() { + s.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(8 * time.Second): + slog.Warn("opencodeSession: close timed out, abandoning wg.Wait") + } close(s.events) return nil } diff --git a/agent/qoder/session.go b/agent/qoder/session.go index 0e3beda..619b6ff 100644 --- a/agent/qoder/session.go +++ b/agent/qoder/session.go @@ -13,6 +13,7 @@ import ( "strings" "sync" "sync/atomic" + "time" "unicode/utf8" "github.com/chenhg5/cc-connect/core" @@ -276,7 +277,16 @@ func (qs *qoderSession) Alive() bool { func (qs *qoderSession) Close() error { qs.alive.Store(false) qs.cancel() - qs.wg.Wait() + done := make(chan struct{}) + go func() { + qs.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(8 * time.Second): + slog.Warn("qoderSession: close timed out, abandoning wg.Wait") + } close(qs.events) return nil } diff --git a/core/engine.go b/core/engine.go index 1e165f1..875419c 100644 --- a/core/engine.go +++ b/core/engine.go @@ -253,14 +253,20 @@ func (e *Engine) Stop() error { e.cancel() e.interactiveMu.Lock() - for key, state := range e.interactiveStates { - if state.agentSession != nil { - state.agentSession.Close() - } - delete(e.interactiveStates, key) + states := make(map[string]*interactiveState, len(e.interactiveStates)) + for k, v := range e.interactiveStates { + states[k] = v + delete(e.interactiveStates, k) } e.interactiveMu.Unlock() + for key, state := range states { + if state.agentSession != nil { + slog.Debug("engine.Stop: closing agent session", "session", key) + state.agentSession.Close() + } + } + var errs []error for _, p := range e.platforms { if err := p.Stop(); err != nil { @@ -561,17 +567,29 @@ func (e *Engine) getOrCreateInteractiveState(sessionKey string, p Platform, repl func (e *Engine) cleanupInteractiveState(sessionKey string) { e.interactiveMu.Lock() - defer e.interactiveMu.Unlock() - state, ok := e.interactiveStates[sessionKey] - if ok && state.agentSession != nil { + delete(e.interactiveStates, sessionKey) + e.interactiveMu.Unlock() + + if ok && state != nil && state.agentSession != nil { + slog.Debug("cleanupInteractiveState: closing agent session", "session", sessionKey) closeStart := time.Now() - state.agentSession.Close() - if elapsed := time.Since(closeStart); elapsed >= slowAgentClose { - slog.Warn("slow agent session close", "elapsed", elapsed, "session", sessionKey) + + done := make(chan struct{}) + go func() { + state.agentSession.Close() + close(done) + }() + + select { + case <-done: + if elapsed := time.Since(closeStart); elapsed >= slowAgentClose { + slog.Warn("slow agent session close", "elapsed", elapsed, "session", sessionKey) + } + case <-time.After(10 * time.Second): + slog.Error("agent session close timed out (10s), abandoning", "session", sessionKey) } } - delete(e.interactiveStates, sessionKey) } func (e *Engine) processInteractiveEvents(state *interactiveState, session *Session, sessionKey string, turnStart time.Time) { @@ -803,7 +821,9 @@ func (e *Engine) handleCommand(p Platform, msg *Message, raw string) bool { } func (e *Engine) cmdNew(p Platform, msg *Message, args []string) { + slog.Info("cmdNew: cleaning up old session", "session_key", msg.SessionKey) e.cleanupInteractiveState(msg.SessionKey) + slog.Info("cmdNew: cleanup done, creating new session", "session_key", msg.SessionKey) name := "session" if len(args) > 0 { name = strings.Join(args, " ") @@ -866,6 +886,7 @@ func (e *Engine) cmdSwitch(p Platform, msg *Message, args []string) { } prefix := strings.TrimSpace(args[0]) + slog.Info("cmdSwitch: listing agent sessions", "session_key", msg.SessionKey) agentSessions, err := e.agent.ListSessions(e.ctx) if err != nil { e.reply(p, msg.ReplyCtx, fmt.Sprintf("❌ %v", err)) @@ -884,7 +905,9 @@ func (e *Engine) cmdSwitch(p Platform, msg *Message, args []string) { return } + slog.Info("cmdSwitch: cleaning up old session", "session_key", msg.SessionKey) e.cleanupInteractiveState(msg.SessionKey) + slog.Info("cmdSwitch: cleanup done", "session_key", msg.SessionKey) session := e.sessions.GetOrCreateActive(msg.SessionKey) session.AgentSessionID = matched.ID