feat(iflow): add configurable tool timeout and reset timer on partial completion

Add toolTimeoutSec field to Agent and Session structs to allow configurable timeout duration.
Reset pending tool timeout timer when partially completing tools.
Increase default timeout from 45s to 180s.
Add tests for timeout behavior and customization.

generated by llmgit

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
chenhg5andClaude committed 2026-03-12 07:04:08 +08:00
1 parent df7d456646
commit dcbfdd672c
3 files changed
+127 -38

No files matched your search

+27 -14
View File
@@ -31,14 +31,15 @@ func init() {
// - "plan": read-only planning mode (--plan)
// - "yolo": auto-approve all tool calls (--yolo)
type Agent struct {
workDir string
model string
mode string
cmd string
providers []core.ProviderConfig
activeIdx int
sessionEnv []string
mu sync.Mutex
workDir string
model string
mode string
cmd string
toolTimeoutSec int
providers []core.ProviderConfig
activeIdx int
sessionEnv []string
mu sync.Mutex
}
func New(opts map[string]any) (core.Agent, error) {
@@ -58,12 +59,23 @@ func New(opts map[string]any) (core.Agent, error) {
return nil, fmt.Errorf("iflow: %q CLI not found in PATH, install with: npm i -g @iflow-ai/iflow-cli", cmd)
}
var toolTimeoutSec int
switch v := opts["tool_timeout_secs"].(type) {
case int64:
toolTimeoutSec = int(v)
case int:
toolTimeoutSec = v
case float64:
toolTimeoutSec = int(v)
}
return &Agent{
workDir: workDir,
model: model,
mode: mode,
cmd: cmd,
activeIdx: -1,
workDir: workDir,
model: model,
mode: mode,
cmd: cmd,
toolTimeoutSec: toolTimeoutSec,
activeIdx: -1,
}, nil
}
@@ -115,6 +127,7 @@ func (a *Agent) StartSession(ctx context.Context, sessionID string) (core.AgentS
mode := a.mode
cmd := a.cmd
workDir := a.workDir
toolTimeoutSec := a.toolTimeoutSec
extraEnv := a.providerEnvLocked()
extraEnv = append(extraEnv, a.sessionEnv...)
if a.activeIdx >= 0 && a.activeIdx < len(a.providers) {
@@ -124,7 +137,7 @@ func (a *Agent) StartSession(ctx context.Context, sessionID string) (core.AgentS
}
a.mu.Unlock()
return newIFlowSession(ctx, cmd, workDir, model, mode, sessionID, extraEnv)
return newIFlowSession(ctx, cmd, workDir, model, mode, sessionID, extraEnv, toolTimeoutSec)
}
func (a *Agent) ListSessions(_ context.Context) ([]core.AgentSessionInfo, error) {
+32 -23
View File
@@ -33,26 +33,27 @@ const (
iflowTranscriptPoll = 200 * time.Millisecond
)
var iflowPendingToolTimeout = 45 * time.Second
var iflowPendingToolTimeout = 180 * time.Second
var iflowPendingToolTimeoutDefaultMode = 6 * time.Second
// iflowSession manages multi-turn conversations with iFlow CLI.
// Each Send() launches a fresh interactive `iflow -i` process inside a PTY,
// then tails the transcript JSONL to recover structured assistant/tool events.
type iflowSession struct {
cmd string
workDir string
model string
mode string
extraEnv []string
events chan core.Event
sessionID atomic.Value // stores string
sentOnce atomic.Bool
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
alive atomic.Bool
turnActive atomic.Bool
cmd string
workDir string
model string
mode string
toolTimeoutSec int
extraEnv []string
events chan core.Event
sessionID atomic.Value // stores string
sentOnce atomic.Bool
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
alive atomic.Bool
turnActive atomic.Bool
}
type iflowTurn struct {
@@ -91,18 +92,19 @@ type iflowToolResult struct {
Output string
}
func newIFlowSession(ctx context.Context, cmd, workDir, model, mode, resumeID string, extraEnv []string) (*iflowSession, error) {
func newIFlowSession(ctx context.Context, cmd, workDir, model, mode, resumeID string, extraEnv []string, toolTimeoutSec int) (*iflowSession, error) {
sessionCtx, cancel := context.WithCancel(ctx)
s := &iflowSession{
cmd: cmd,
workDir: workDir,
model: model,
mode: mode,
extraEnv: extraEnv,
events: make(chan core.Event, 64),
ctx: sessionCtx,
cancel: cancel,
cmd: cmd,
workDir: workDir,
model: model,
mode: mode,
toolTimeoutSec: toolTimeoutSec,
extraEnv: extraEnv,
events: make(chan core.Event, 64),
ctx: sessionCtx,
cancel: cancel,
}
s.alive.Store(true)
@@ -301,6 +303,9 @@ func (s *iflowSession) watchTranscript(turn *iflowTurn) {
}
func (s *iflowSession) pendingToolTimeout() time.Duration {
if s.toolTimeoutSec > 0 {
return time.Duration(s.toolTimeoutSec) * time.Second
}
if strings.EqualFold(s.mode, "default") {
return iflowPendingToolTimeoutDefaultMode
}
@@ -671,6 +676,7 @@ func (t *iflowTurn) completeTools(results []iflowToolResult) bool {
t.mu.Lock()
defer t.mu.Unlock()
changed := false
for _, result := range results {
if _, ok := t.doneToolIDs[result.ID]; ok {
continue
@@ -681,10 +687,13 @@ func (t *iflowTurn) completeTools(results []iflowToolResult) bool {
if result.Output != "" {
t.toolFallback = append(t.toolFallback, result.Output)
}
changed = true
}
if len(t.pendingToolIDs) == 0 && t.pendingTimer != nil {
t.pendingTimer.Stop()
t.pendingTimer = nil
} else if changed && len(t.pendingToolIDs) > 0 && t.pendingTimer != nil {
t.pendingTimer.Reset(t.pendingTimeout)
}
if len(results) == 0 || len(t.pendingToolIDs) > 0 {
return false
+68 -1
View File
@@ -259,6 +259,7 @@ func TestIFlowTurnPendingToolTimeoutReleasesTurn(t *testing.T) {
turn := &iflowTurn{
cancel: func() {},
pendingTimeout: iflowPendingToolTimeout,
pendingToolIDs: make(map[string]struct{}),
pendingTools: make(map[string]iflowToolUse),
seenToolIDs: make(map[string]struct{}),
@@ -288,6 +289,72 @@ func TestIFlowTurnPendingToolTimeoutReleasesTurn(t *testing.T) {
t.Fatal("pending tool timeout did not release turn")
}
func TestIFlowTurnTimerResetsOnPartialToolCompletion(t *testing.T) {
timeout := 100 * time.Millisecond
cancelled := false
turn := &iflowTurn{
cancel: func() { cancelled = true },
pendingTimeout: timeout,
pendingToolIDs: make(map[string]struct{}),
pendingTools: make(map[string]iflowToolUse),
seenToolIDs: make(map[string]struct{}),
doneToolIDs: make(map[string]struct{}),
}
defer turn.stopResultTimer()
turn.addPendingTools([]iflowToolUse{
{ID: "write_file:1", Name: "write_file"},
{ID: "run_shell_command:2", Name: "run_shell_command"},
})
// Wait 70ms (>50% of timeout), then complete one tool
time.Sleep(70 * time.Millisecond)
turn.completeTools([]iflowToolResult{{ID: "write_file:1", Output: "ok"}})
// Timer was reset — wait another 70ms; should NOT have timed out yet
time.Sleep(70 * time.Millisecond)
if turn.readyForResult() {
t.Fatal("timer should have been reset; should not be ready yet")
}
// Now wait for the full reset timeout to expire
time.Sleep(50 * time.Millisecond)
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
if turn.readyForResult() {
if !cancelled {
t.Fatal("expected cancel to be called")
}
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatal("timer did not fire after reset")
}
func TestIFlowSessionCustomToolTimeout(t *testing.T) {
sess, err := newIFlowSession(context.Background(), "echo", "/tmp", "", "yolo", "", nil, 300)
if err != nil {
t.Fatalf("newIFlowSession: %v", err)
}
defer sess.Close()
if got := sess.pendingToolTimeout(); got != 300*time.Second {
t.Fatalf("pendingToolTimeout = %v, want 300s", got)
}
}
func TestIFlowSessionDefaultToolTimeout(t *testing.T) {
sess, err := newIFlowSession(context.Background(), "echo", "/tmp", "", "yolo", "", nil, 0)
if err != nil {
t.Fatalf("newIFlowSession: %v", err)
}
defer sess.Close()
if got := sess.pendingToolTimeout(); got != iflowPendingToolTimeout {
t.Fatalf("pendingToolTimeout = %v, want %v", got, iflowPendingToolTimeout)
}
}
func TestIFlowSessionPendingToolTimeoutClearsBusyState(t *testing.T) {
oldTimeout := iflowPendingToolTimeout
oldDefaultTimeout := iflowPendingToolTimeoutDefaultMode
@@ -324,7 +391,7 @@ while :; do sleep 1; done
t.Fatalf("WriteFile fake iflow: %v", err)
}
sess, err := newIFlowSession(context.Background(), cmdPath, workDir, "", "default", "", nil)
sess, err := newIFlowSession(context.Background(), cmdPath, workDir, "", "default", "", nil, 0)
if err != nil {
t.Fatalf("newIFlowSession: %v", err)
}