mirror of
https://github.com/coder/coder.git
synced 2026-10-31 02:43:21 +08:00
feat(coderd/x/chatd): bill interrupted tool batches on cancellation rows (#28362)
**Stack**: #28359 → #28360 → #28361 → **#28362 (this PR)** ## Summary Interrupted local tool batches now bill the partial window they actually executed, recorded on the same dedicated usage record completed batches use. Completes the runtime billing work in this stack: completed batches bill via #28360, interrupted batches via this PR. Last of a four-PR stack splitting #28211. Stacked on #28361. Closes CODAGT-928 (https://linear.app/codercom/issue/CODAGT-928/track-local-tool-execution-for-agent-runtime). ## Problem Interrupting a batch previously billed only the model window; work the tools already performed was lost. ## Fix Just before closing the buffer episode, the interrupt task reads the live tool stamps (`ToolCompletions`) alongside `ModelInvokedAt`. The interruption commit appends the shared `batchUsageMessage` record (same shape as #28360: model-only tool-role row with a `billed_ms`/`billed_calls` audit payload) billing the union of the started, billable calls' partial execution intervals: completed calls end at their recorded completion, running calls at the interrupt instant. Synthesized cancellation rows stay free of batch-level runtime. Unstarted calls, rejected calls, and sub-agent orchestration tools bill nothing and append no record. Batches with no live attempt (crash recovery, state promotion) bill nothing, consistent with generation losing in-flight runtime on a crash. Usage-record billing reuses the same interrupt instant as the model window instead of taking a second clock reading inside the transaction, so database delay cannot inflate a running call's window. An earlier revision carried a retry-stable episode snapshot so an interrupt retry after buffer eviction (a >15s stall) would not lose billing; it was dropped as not worth the complexity (https://github.com/coder/coder/pull/28211#discussion_r3812778418). Such a retry now bills the batch at zero, matching the existing crash-recovery precedent. Known gaps left out of scope: retry backoff and errored-attempt time stay unbilled (no persistence home for error turns), and the advisor tool's nested model call now bills as tool time by design.
This commit is contained in:
Jaayden Halko
+40
-2
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatdebug"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatloop"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer"
|
||||
@@ -261,6 +262,7 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt
|
||||
GenerationAttempt: chat.GenerationAttempt,
|
||||
}
|
||||
modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key)
|
||||
toolCompletions := s.opts.MessagePartBuffer.ToolCompletions(key)
|
||||
if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err())
|
||||
@@ -302,7 +304,9 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt
|
||||
return xerrors.Errorf("load chat for task: %w", err)
|
||||
}
|
||||
messages := partialMessages
|
||||
committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"))
|
||||
// Reuse the captured interrupt instant so database delay does not
|
||||
// inflate billing.
|
||||
committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, interruptedAt, toolCompletions)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("committed pending local tool cancellation messages: %w", err)
|
||||
}
|
||||
@@ -685,6 +689,7 @@ func committedPendingLocalToolCancellationMessages(
|
||||
store database.Store,
|
||||
chat database.Chat,
|
||||
interruptedAt time.Time,
|
||||
toolCompletions map[int]messagepartbuffer.ToolCompletion,
|
||||
) ([]chatstate.Message, error) {
|
||||
messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
|
||||
ChatID: chat.ID,
|
||||
@@ -700,8 +705,9 @@ func committedPendingLocalToolCancellationMessages(
|
||||
if len(localCalls) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var intervals []chatloop.BilledInterval
|
||||
result := make([]chatstate.Message, 0, len(localCalls))
|
||||
for _, call := range localCalls {
|
||||
for i, call := range localCalls {
|
||||
payload, err := json.Marshal(map[string]string{"error": interruptedToolResultErrorMessage})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("marshal interrupted tool result: %w", err)
|
||||
@@ -721,6 +727,38 @@ func committedPendingLocalToolCancellationMessages(
|
||||
ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil},
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
})
|
||||
if unbilledSubagentToolNames[call.ToolName] {
|
||||
continue
|
||||
}
|
||||
// Bill only started calls. Completed calls end at completion and
|
||||
// running calls end at the interrupt.
|
||||
occurrence, ok := toolCompletions[i]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
start := occurrence.StartedAt
|
||||
if start.IsZero() {
|
||||
continue
|
||||
}
|
||||
end := occurrence.CompletedAt
|
||||
if end.IsZero() {
|
||||
end = interruptedAt
|
||||
}
|
||||
intervals = append(intervals, chatloop.BilledInterval{Start: start, End: end})
|
||||
}
|
||||
// Bill the interval union once on a dedicated usage record so
|
||||
// cancellation rows stay free of batch-level runtime.
|
||||
stamp, ok, err := batchUsageMessage(
|
||||
chat.LastModelConfigID,
|
||||
chatprompt.CurrentContentVersion,
|
||||
chatloop.BilledIntervalsDuration(intervals),
|
||||
len(intervals),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if ok {
|
||||
result = append(result, stamp)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -475,6 +475,375 @@ func TestInterruptTask_PartialAssistantWithoutModelInvocationHasNoRuntime(t *tes
|
||||
require.False(t, assistant.RuntimeMs.Valid)
|
||||
}
|
||||
|
||||
type interruptedBatch struct {
|
||||
chat database.Chat
|
||||
starter *taskStarter
|
||||
clock *quartz.Mock
|
||||
key messagepartbuffer.Key
|
||||
workerID uuid.UUID
|
||||
runnerID uuid.UUID
|
||||
}
|
||||
|
||||
func interruptedBatchFixture(
|
||||
t *testing.T,
|
||||
f *taskTestFixture,
|
||||
calls []codersdk.ChatMessagePart,
|
||||
) interruptedBatch {
|
||||
t.Helper()
|
||||
chat := f.createRunningChat(t)
|
||||
raw, err := chatprompt.MarshalParts(calls)
|
||||
require.NoError(t, err)
|
||||
machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID)
|
||||
require.NoError(t, machine.Update(testutil.Context(t, testutil.WaitShort), func(tx *chatstate.Tx, _ database.Store) error {
|
||||
_, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{{
|
||||
Role: database.ChatMessageRoleAssistant,
|
||||
Content: raw,
|
||||
Visibility: database.ChatMessageVisibilityBoth,
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
ModelConfigID: uuid.NullUUID{UUID: f.model.ID, Valid: true},
|
||||
}}})
|
||||
return err
|
||||
}))
|
||||
workerID := uuid.New()
|
||||
runnerID := uuid.New()
|
||||
acquired := f.acquireChat(t, chat.ID, workerID, runnerID)
|
||||
recorder := newTaskSideEffectRecorder()
|
||||
clock := quartz.NewMock(t)
|
||||
starter := newTestTaskStarterWithClock(t, f, recorder, clock)
|
||||
key := messagepartbuffer.Key{
|
||||
ChatID: chat.ID,
|
||||
HistoryVersion: acquired.HistoryVersion,
|
||||
GenerationAttempt: acquired.GenerationAttempt,
|
||||
}
|
||||
require.NoError(t, starter.opts.MessagePartBuffer.CreateEpisode(key))
|
||||
return interruptedBatch{
|
||||
chat: chat,
|
||||
starter: starter,
|
||||
clock: clock,
|
||||
key: key,
|
||||
workerID: workerID,
|
||||
runnerID: runnerID,
|
||||
}
|
||||
}
|
||||
|
||||
func (b interruptedBatch) interrupt(t *testing.T, f *taskTestFixture) []database.ChatMessage {
|
||||
t.Helper()
|
||||
interrupting := f.interruptChat(t, b.chat.ID)
|
||||
err := b.starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{
|
||||
ChatID: b.chat.ID,
|
||||
WorkerID: b.workerID,
|
||||
RunnerID: b.runnerID,
|
||||
HistoryVersion: interrupting.HistoryVersion,
|
||||
GenerationAttempt: interrupting.GenerationAttempt,
|
||||
Status: database.ChatStatusInterrupting,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: b.chat.ID})
|
||||
require.NoError(t, err)
|
||||
return messages
|
||||
}
|
||||
|
||||
func findToolResultMessage(t *testing.T, messages []database.ChatMessage, toolCallID string) database.ChatMessage {
|
||||
t.Helper()
|
||||
for _, msg := range messages {
|
||||
if msg.Role != database.ChatMessageRoleTool {
|
||||
continue
|
||||
}
|
||||
parts, err := chatprompt.ParseContent(msg)
|
||||
require.NoError(t, err)
|
||||
for _, part := range parts {
|
||||
if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolCallID == toolCallID {
|
||||
return msg
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Fatalf("no tool result message for call %s", toolCallID)
|
||||
return database.ChatMessage{}
|
||||
}
|
||||
|
||||
// batchUsageRecords returns the model-only tool batch usage rows for the
|
||||
// chat. Cancellation rows are user-visible, so the dedicated usage record
|
||||
// only appears in the model-visibility history.
|
||||
func batchUsageRecords(t *testing.T, f *taskTestFixture, chatID uuid.UUID) []database.ChatMessage {
|
||||
t.Helper()
|
||||
rows, err := f.db.GetChatMessagesForPromptByChatID(testutil.Context(t, testutil.WaitShort), chatID)
|
||||
require.NoError(t, err)
|
||||
var records []database.ChatMessage
|
||||
for _, msg := range rows {
|
||||
if msg.Role != database.ChatMessageRoleTool || msg.Visibility != database.ChatMessageVisibilityModel {
|
||||
continue
|
||||
}
|
||||
records = append(records, msg)
|
||||
}
|
||||
return records
|
||||
}
|
||||
|
||||
func requireSingleBatchUsageRecord(t *testing.T, f *taskTestFixture, chatID uuid.UUID, billedMs int64, billedCalls int) {
|
||||
t.Helper()
|
||||
records := batchUsageRecords(t, f, chatID)
|
||||
require.Len(t, records, 1)
|
||||
require.Equal(t, sql.NullInt64{Int64: billedMs, Valid: true}, records[0].RuntimeMs)
|
||||
parts, err := chatprompt.ParseContent(records[0])
|
||||
require.NoError(t, err)
|
||||
require.Len(t, parts, 1)
|
||||
require.Equal(t, toolBatchUsagePartType, parts[0].Type)
|
||||
require.JSONEq(t,
|
||||
fmt.Sprintf(`{"billed_ms":%d,"billed_calls":%d}`, billedMs, billedCalls),
|
||||
string(parts[0].Result),
|
||||
)
|
||||
}
|
||||
|
||||
func TestInterruptTask_ToolBatchBillsPartialWindowOnUsageRecord(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
execCallID := "call_" + uuid.NewString()
|
||||
waitCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
buffer := batch.starter.opts.MessagePartBuffer
|
||||
|
||||
// Keep advances below the buffer's 15-second cleanup tick.
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now()))
|
||||
// Record a live completion without publishing a tool result.
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, batch.clock.Now()))
|
||||
batch.clock.Advance(5 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
execRow := findToolResultMessage(t, messages, execCallID)
|
||||
require.False(t, execRow.RuntimeMs.Valid)
|
||||
waitRow := findToolResultMessage(t, messages, waitCallID)
|
||||
require.False(t, waitRow.RuntimeMs.Valid)
|
||||
requireSingleBatchUsageRecord(t, f, batch.chat.ID, 3_000, 1)
|
||||
}
|
||||
|
||||
func TestInterruptTask_RunningBilledToolBillsUpToInterrupt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
execCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
batch.clock.Advance(7 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
execRow := findToolResultMessage(t, messages, execCallID)
|
||||
require.False(t, execRow.RuntimeMs.Valid)
|
||||
requireSingleBatchUsageRecord(t, f, batch.chat.ID, 7_000, 1)
|
||||
}
|
||||
|
||||
func TestInterruptTask_UnbilledOnlyBatchBillsNothingOnInterrupt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
waitCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
batch.clock.Advance(10 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
waitRow := findToolResultMessage(t, messages, waitCallID)
|
||||
require.False(t, waitRow.RuntimeMs.Valid)
|
||||
require.Empty(t, batchUsageRecords(t, f, batch.chat.ID))
|
||||
}
|
||||
|
||||
func TestInterruptTask_UnstartedSerialCallBillsNothingOnInterrupt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
waitCallID := "call_" + uuid.NewString()
|
||||
serialCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: serialCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
buffer := batch.starter.opts.MessagePartBuffer
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
batch.clock.Advance(10 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
for _, msg := range messages {
|
||||
if msg.Role == database.ChatMessageRoleTool {
|
||||
require.False(t, msg.RuntimeMs.Valid,
|
||||
"no cancellation row may bill: the billed call never began executing")
|
||||
}
|
||||
}
|
||||
require.Empty(t, batchUsageRecords(t, f, batch.chat.ID),
|
||||
"no usage record may bill: the billed call never began executing")
|
||||
}
|
||||
|
||||
func TestInterruptTask_StartedSerialCallBillsFromItsOwnStart(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
execCallID := "call_" + uuid.NewString()
|
||||
waitCallID := "call_" + uuid.NewString()
|
||||
serialCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: serialCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
buffer := batch.starter.opts.MessagePartBuffer
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now()))
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, batch.clock.Now()))
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, batch.clock.Now()))
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 2, batch.clock.Now()))
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
serialRow := findToolResultMessage(t, messages, serialCallID)
|
||||
require.False(t, serialRow.RuntimeMs.Valid)
|
||||
execRow := findToolResultMessage(t, messages, execCallID)
|
||||
require.False(t, execRow.RuntimeMs.Valid)
|
||||
waitRow := findToolResultMessage(t, messages, waitCallID)
|
||||
require.False(t, waitRow.RuntimeMs.Valid)
|
||||
// Completed execute [2s,5s] plus the serial call's own window
|
||||
// [8s,10s]: the gap and the unbilled wait_agent do not count.
|
||||
requireSingleBatchUsageRecord(t, f, batch.chat.ID, 5_000, 2)
|
||||
}
|
||||
|
||||
func TestInterruptTask_DuplicateCallIDsKeepOccurrenceStates(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
dupCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
buffer := batch.starter.opts.MessagePartBuffer
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now()))
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolCompletion(batch.key, 0, batch.clock.Now()))
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
for _, msg := range messages {
|
||||
if msg.Role == database.ChatMessageRoleTool {
|
||||
require.False(t, msg.RuntimeMs.Valid)
|
||||
}
|
||||
}
|
||||
// The still-running occurrence bills the full window exactly once.
|
||||
requireSingleBatchUsageRecord(t, f, batch.chat.ID, 6_000, 2)
|
||||
}
|
||||
|
||||
func TestInterruptTask_DuplicateCallIDCompletionStampsOwnOccurrence(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
dupCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
buffer := batch.starter.opts.MessagePartBuffer
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 0, batch.clock.Now()))
|
||||
require.NoError(t, buffer.RecordToolStart(batch.key, 1, batch.clock.Now()))
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
require.NoError(t, buffer.RecordToolCompletion(batch.key, 1, batch.clock.Now()))
|
||||
batch.clock.Advance(3 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
for _, msg := range messages {
|
||||
if msg.Role == database.ChatMessageRoleTool {
|
||||
require.False(t, msg.RuntimeMs.Valid)
|
||||
}
|
||||
}
|
||||
// The running execute occurrence bills to the interrupt; wait_agent's
|
||||
// early completion must not end it at 3s.
|
||||
requireSingleBatchUsageRecord(t, f, batch.chat.ID, 6_000, 1)
|
||||
}
|
||||
|
||||
func TestInterruptTask_RejectedDuplicateIDDoesNotStealDispatchedOccurrence(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
dupCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: dupCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 1, batch.clock.Now()))
|
||||
batch.clock.Advance(10 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
for _, msg := range messages {
|
||||
if msg.Role == database.ChatMessageRoleTool {
|
||||
require.False(t, msg.RuntimeMs.Valid)
|
||||
}
|
||||
}
|
||||
require.Empty(t, batchUsageRecords(t, f, batch.chat.ID),
|
||||
"no usage record may bill: only the unbilled wait_agent occurrence ran")
|
||||
}
|
||||
|
||||
func TestInterruptTask_RejectedCallBillsNothingOnInterrupt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
execCallID := "call_" + uuid.NewString()
|
||||
waitCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: waitCallID, ToolName: "wait_agent", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
|
||||
batch.clock.Advance(2 * time.Second)
|
||||
require.NoError(t, batch.starter.opts.MessagePartBuffer.RecordToolStart(batch.key, 1, batch.clock.Now()))
|
||||
batch.clock.Advance(10 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
execRow := findToolResultMessage(t, messages, execCallID)
|
||||
require.False(t, execRow.RuntimeMs.Valid)
|
||||
waitRow := findToolResultMessage(t, messages, waitCallID)
|
||||
require.False(t, waitRow.RuntimeMs.Valid)
|
||||
require.Empty(t, batchUsageRecords(t, f, batch.chat.ID))
|
||||
}
|
||||
|
||||
func TestInterruptTask_ToolCancellationWithoutLiveBatchHasNoRuntime(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
f := newTaskTestFixture(t)
|
||||
execCallID := "call_" + uuid.NewString()
|
||||
batch := interruptedBatchFixture(t, f, []codersdk.ChatMessagePart{
|
||||
{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: execCallID, ToolName: "execute", Args: json.RawMessage(`{}`)},
|
||||
})
|
||||
|
||||
batch.clock.Advance(10 * time.Second)
|
||||
|
||||
messages := batch.interrupt(t, f)
|
||||
execRow := findToolResultMessage(t, messages, execCallID)
|
||||
require.False(t, execRow.RuntimeMs.Valid)
|
||||
require.Empty(t, batchUsageRecords(t, f, batch.chat.ID))
|
||||
}
|
||||
|
||||
func TestRequiresActionTimeout_ExpiredCancelsOnly(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user