diff --git a/agent/iflow/iflow.go b/agent/iflow/iflow.go index 885447d..9b56787 100644 --- a/agent/iflow/iflow.go +++ b/agent/iflow/iflow.go @@ -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) { diff --git a/agent/iflow/session.go b/agent/iflow/session.go index aa88eaa..227dac8 100644 --- a/agent/iflow/session.go +++ b/agent/iflow/session.go @@ -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 diff --git a/agent/iflow/session_test.go b/agent/iflow/session_test.go index 3c9fd26..e21a561 100644 --- a/agent/iflow/session_test.go +++ b/agent/iflow/session_test.go @@ -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) }