mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-16 05:00:53 +08:00
Compare commits
12
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
60728092e1 | ||
|
|
ff599c65d3 | ||
|
|
ce096cdba1 | ||
|
|
50f6d07640 | ||
|
|
e6ab0f34a7 | ||
|
|
5ed28cde9f | ||
|
|
773fa9183a | ||
|
|
1cbab7cc39 | ||
|
|
bc75b5dc43 | ||
|
|
cddcd98f3b | ||
|
|
d42988b6a8 | ||
|
|
3975e58ce0 |
+8
-2
@@ -169,9 +169,12 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin
|
||||
if c.guard != nil {
|
||||
opts.Hooks = c.guard.Hooks()
|
||||
}
|
||||
// The session id is the provider-side identity (session header, usage metering)
|
||||
// with or without persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = sessionID
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if c.tx != nil { // persist raw human↔AI conversation; one accumulating file per thread
|
||||
opts.Transcript = c.tx
|
||||
opts.SessionID = sessionID
|
||||
}
|
||||
// experimental feature: when enabled, noa takes over context compression (archive centralized under <workDir>/noa/<SessionID>, persistent).
|
||||
enableNoa(&opts, c.noaEnabledFn, c.workDir, "chat-"+sessionID, noaWarn("chat-"+sessionID))
|
||||
@@ -180,8 +183,11 @@ func (c *ChatAgent) Chat(ctx context.Context, agentKey, sessionID, message strin
|
||||
defer s.Close()
|
||||
// reload prior conversation so the agent has context across turns (each Chat is
|
||||
// a fresh session). First turn: no file yet → Resume loads nothing and proceeds.
|
||||
// A corrupt or unreadable transcript stops the turn before any model or tool call.
|
||||
if c.tx != nil {
|
||||
_ = s.Resume(sessionID)
|
||||
if err := resumeSession(s, opts.SessionID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
// re-unlock skill-gated MCPs from prior Skill() calls in the reloaded history so
|
||||
// revealed tools stay callable across the fresh session.
|
||||
|
||||
+13
-7
@@ -169,14 +169,17 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, mainSeg int, as *db.
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
// Segment 0 keeps the legacy "exp%d-main" name so existing transcripts still
|
||||
// load; each new session (seg>=1) gets its own file for a clean context. The id
|
||||
// is the provider-side identity (session header, usage metering) with or without
|
||||
// persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main", ts.ID())
|
||||
if mainSeg > 0 {
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main-s%d", ts.ID(), mainSeg)
|
||||
}
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if m.tx != nil { // persist raw human↔AI conversation; one accumulating file per segment
|
||||
opts.Transcript = m.tx
|
||||
// Segment 0 keeps the legacy "exp%d-main" name so existing transcripts still
|
||||
// load; each new session (seg>=1) gets its own file for a clean context.
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main", ts.ID())
|
||||
if mainSeg > 0 {
|
||||
opts.SessionID = fmt.Sprintf("exp%d-main-s%d", ts.ID(), mainSeg)
|
||||
}
|
||||
}
|
||||
// experimental feature: when enabled, noa takes over context compression (archive centralized under <workDir>/noa/<SessionID>, persistent).
|
||||
// the session id follows the same rule as the transcript (segment-aware), keeping archive and resume aligned.
|
||||
@@ -191,8 +194,11 @@ func (m *MainAgent) Chat(ctx context.Context, taskID int64, mainSeg int, as *db.
|
||||
// reload the prior conversation from the transcript so the agent has context
|
||||
// across turns (each Chat is a fresh session; without this it can't see earlier
|
||||
// messages). First turn: no file yet → Resume loads nothing and proceeds.
|
||||
// A corrupt or unreadable transcript stops the turn before any model or tool call.
|
||||
if m.tx != nil {
|
||||
_ = s.Resume(opts.SessionID)
|
||||
if err := resumeSession(s, opts.SessionID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
// C2: this session is fresh each turn; re-unlock skill-gated MCPs from prior
|
||||
// Skill() calls in the reloaded history so revealed tools stay callable.
|
||||
|
||||
+4
-1
@@ -567,9 +567,12 @@ func (p *Planner) Plan(ctx context.Context, taskID int64, as *db.AssetStore, ts
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
// The session id is the provider-side identity (session header, usage metering)
|
||||
// with or without persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if p.tx != nil { // persist raw LLM conversation; one accumulating file per task's planner
|
||||
opts.Transcript = p.tx
|
||||
opts.SessionID = fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
}
|
||||
// Experimental: when on, noa takes over context compression (archives persist under <workDir>/noa/<SessionID>).
|
||||
noaSession := fmt.Sprintf("exp%d-planner", ts.ID())
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
)
|
||||
|
||||
// resumeSession reloads the persisted conversation of session id into s. Norma's
|
||||
// transcript.Store.Load treats a missing file as a fresh session (no records, no
|
||||
// error); any other failure, such as a corrupt line or an unreadable path, is
|
||||
// returned so the caller stops before the first model or tool call instead of
|
||||
// appending a new conversation to the damaged file.
|
||||
func resumeSession(s *agentcore.Session, id string) error {
|
||||
if err := s.Resume(id); err != nil {
|
||||
return fmt.Errorf("restore Norma session %q: %w", id, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/transcript"
|
||||
)
|
||||
|
||||
// sessionProbe is a provider that answers every call with one text turn and
|
||||
// records the session id it saw on ctx plus the request messages.
|
||||
type sessionProbe struct {
|
||||
mu sync.Mutex
|
||||
ids []string
|
||||
msgs [][]llm.Message
|
||||
}
|
||||
|
||||
func (p *sessionProbe) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
p.mu.Lock()
|
||||
p.ids = append(p.ids, transcript.SessionIDFrom(ctx))
|
||||
p.msgs = append(p.msgs, req.Messages)
|
||||
p.mu.Unlock()
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
if yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "done"}, nil) {
|
||||
yield(llm.StreamEvent{Type: llm.SEMessageDelta, StopReason: "end_turn"}, nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *sessionProbe) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
acc := llm.NewAccumulator()
|
||||
for ev, err := range p.Stream(ctx, req) {
|
||||
if err != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, err
|
||||
}
|
||||
acc.Add(ev)
|
||||
}
|
||||
return acc.Message(), acc.StopReason, acc.Usage, nil
|
||||
}
|
||||
|
||||
func (p *sessionProbe) seen() []string {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return append([]string(nil), p.ids...)
|
||||
}
|
||||
|
||||
// jsonlFiles lists every .jsonl file under dir.
|
||||
func jsonlFiles(t *testing.T, dir string) []string {
|
||||
t.Helper()
|
||||
var out []string
|
||||
err := filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !d.IsDir() && strings.HasSuffix(path, ".jsonl") {
|
||||
out = append(out, path)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// damageTranscript makes the main transcript of id corrupt (a non-JSON line) or
|
||||
// unreadable (a directory). It returns the bytes a corrupt file must keep.
|
||||
func damageTranscript(t *testing.T, store *transcript.Store, id string, unreadable bool) string {
|
||||
t.Helper()
|
||||
path := store.MainPath(id)
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if unreadable {
|
||||
if err := os.Mkdir(path, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
const original = "{corrupt transcript\n"
|
||||
if err := os.WriteFile(path, []byte(original), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return original
|
||||
}
|
||||
|
||||
// assertTranscriptUntouched checks that a damaged transcript was neither
|
||||
// rewritten nor appended to.
|
||||
func assertTranscriptUntouched(t *testing.T, store *transcript.Store, id, original string, unreadable bool) {
|
||||
t.Helper()
|
||||
path := store.MainPath(id)
|
||||
if unreadable {
|
||||
entries, err := os.ReadDir(path)
|
||||
if err != nil || len(entries) != 0 {
|
||||
t.Fatalf("unreadable transcript changed: entries=%v err=%v", entries, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil || string(b) != original {
|
||||
t.Fatalf("damaged transcript changed: %q err=%v", b, err)
|
||||
}
|
||||
}
|
||||
|
||||
func assertRestoreError(t *testing.T, err error, calls int) {
|
||||
t.Helper()
|
||||
if err == nil || !strings.Contains(err.Error(), "restore Norma session") || calls != 0 {
|
||||
t.Fatalf("err=%v provider calls=%d", err, calls)
|
||||
}
|
||||
}
|
||||
|
||||
var damageCases = []struct {
|
||||
name string
|
||||
unreadable bool
|
||||
}{{"corrupt", false}, {"unreadable-directory", true}}
|
||||
|
||||
func TestChatSessionIdentityWithoutTranscript(t *testing.T) {
|
||||
p := &sessionProbe{}
|
||||
workDir := t.TempDir()
|
||||
c := NewChatAgent(p, "test-model", workDir, nil, 100000)
|
||||
for range 2 {
|
||||
if _, err := c.Chat(t.Context(), "assistant", "conv-stable", "hello", 1, 0, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if ids := p.seen(); len(ids) != 2 || ids[0] != "conv-stable" || ids[1] != "conv-stable" {
|
||||
t.Fatalf("provider session ids: %v", ids)
|
||||
}
|
||||
if files := jsonlFiles(t, workDir); len(files) != 0 {
|
||||
t.Fatalf("transcript written without persistence: %v", files)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatMissingTranscriptStartsFresh(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
p := &sessionProbe{}
|
||||
c := NewChatAgent(p, "test-model", t.TempDir(), store, 100000)
|
||||
if _, err := c.Chat(t.Context(), "assistant", "conv-new", "hello", 1, 0, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ids := p.seen(); len(ids) != 1 || ids[0] != "conv-new" {
|
||||
t.Fatalf("provider session ids: %v", ids)
|
||||
}
|
||||
recs, err := store.Load(store.MainPath("conv-new"))
|
||||
if err != nil || len(recs) == 0 {
|
||||
t.Fatalf("fresh transcript not written: %d records, err=%v", len(recs), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatRefusesDamagedTranscript(t *testing.T) {
|
||||
for _, tc := range damageCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
original := damageTranscript(t, store, "conv-broken", tc.unreadable)
|
||||
p := &sessionProbe{}
|
||||
c := NewChatAgent(p, "test-model", t.TempDir(), store, 100000)
|
||||
_, err := c.Chat(t.Context(), "assistant", "conv-broken", "hello", 1, 0, false, nil)
|
||||
assertRestoreError(t, err, len(p.seen()))
|
||||
assertTranscriptUntouched(t, store, "conv-broken", original, tc.unreadable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResumesSignedThinkingHistory(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
w := store.NewWriter("conv-signed", "")
|
||||
w.RecordMessage(llm.UserText("earlier question"), llm.Usage{})
|
||||
w.RecordMessage(llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{
|
||||
{Type: llm.BlockThinking, Thinking: "reasoning", Signature: "signed"},
|
||||
{Type: llm.BlockRedactedThinking, Data: "opaque"},
|
||||
{Type: llm.BlockText, Text: "earlier answer"},
|
||||
}}, llm.Usage{})
|
||||
if err := w.Err(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p := &sessionProbe{}
|
||||
c := NewChatAgent(p, "test-model", t.TempDir(), store, 100000)
|
||||
if _, err := c.Chat(t.Context(), "assistant", "conv-signed", "next question", 1, 0, false, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if len(p.msgs) != 1 {
|
||||
t.Fatalf("provider calls: %d", len(p.msgs))
|
||||
}
|
||||
var thinking, redacted bool
|
||||
for _, m := range p.msgs[0] {
|
||||
for _, b := range m.Content {
|
||||
thinking = thinking || (b.Type == llm.BlockThinking && b.Thinking == "reasoning" && b.Signature == "signed")
|
||||
redacted = redacted || (b.Type == llm.BlockRedactedThinking && b.Data == "opaque")
|
||||
}
|
||||
}
|
||||
if !thinking || !redacted {
|
||||
t.Fatalf("resumed history lost signed thinking=%v redacted=%v: %+v", thinking, redacted, p.msgs[0])
|
||||
}
|
||||
}
|
||||
|
||||
// sessionFixture is one exploration with one intent on the test database.
|
||||
type sessionFixture struct {
|
||||
expID int64
|
||||
intentID int64
|
||||
ts *db.ExplorationStore
|
||||
intent *db.Node
|
||||
}
|
||||
|
||||
func newSessionFixture(t *testing.T) sessionFixture {
|
||||
t.Helper()
|
||||
d := testDB(t)
|
||||
t.Cleanup(func() { _ = d.Close() })
|
||||
expID, err := d.CreateExploration("session identity probe", "keep the session id stable")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = d.Exec(`DELETE FROM explorations WHERE id=$1`, expID) })
|
||||
ts := d.Exploration(expID)
|
||||
payload := map[string]any{"summary": "session probe"}
|
||||
intentID, err := ts.AddNode("intent", payload, 0, "running", "planner", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
return sessionFixture{expID: expID, intentID: intentID, ts: ts, intent: &db.Node{ID: intentID, Payload: raw}}
|
||||
}
|
||||
|
||||
// run drives one agent kind once and returns the session id it should use.
|
||||
func (f sessionFixture) run(t *testing.T, kind string, p *sessionProbe, tx *transcript.Store) (string, error) {
|
||||
t.Helper()
|
||||
workDir := t.TempDir()
|
||||
switch kind {
|
||||
case "worker":
|
||||
_, _, err := NewWorker(p, "test-model", workDir, tx, 0, 3).Execute(t.Context(), "work#1", f.expID, nil, f.ts, f.intent, passHooks{}, nil, nil, nil)
|
||||
return WorkerSessionID(f.expID, f.intentID), err
|
||||
case "planner":
|
||||
_, _, err := NewPlanner(p, "test-model", workDir, tx, 0, 3).Plan(t.Context(), f.expID, nil, f.ts, "probe goal", nil, passHooks{}, nil)
|
||||
return fmt.Sprintf("exp%d-planner", f.expID), err
|
||||
case "mainagent":
|
||||
_, err := NewMainAgent(p, "test-model", workDir, tx, 0, 3).Chat(t.Context(), f.expID, 0, nil, f.ts, "probe goal", "status?", passHooks{}, nil, nil, nil, nil, nil)
|
||||
return fmt.Sprintf("exp%d-main", f.expID), err
|
||||
case "mainagent-seg":
|
||||
_, err := NewMainAgent(p, "test-model", workDir, tx, 0, 3).Chat(t.Context(), f.expID, 2, nil, f.ts, "probe goal", "status?", passHooks{}, nil, nil, nil, nil, nil)
|
||||
return fmt.Sprintf("exp%d-main-s2", f.expID), err
|
||||
}
|
||||
t.Fatalf("unknown agent kind %q", kind)
|
||||
return "", nil
|
||||
}
|
||||
|
||||
func TestAgentSessionIdentityWithAndWithoutTranscript(t *testing.T) {
|
||||
f := newSessionFixture(t)
|
||||
for _, kind := range []string{"worker", "planner", "mainagent", "mainagent-seg"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
bare := &sessionProbe{}
|
||||
want, err := f.run(t, kind, bare, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
persisted := &sessionProbe{}
|
||||
if _, err := f.run(t, kind, persisted, store); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, ids := range [][]string{bare.seen(), persisted.seen()} {
|
||||
if len(ids) == 0 {
|
||||
t.Fatal("provider never called")
|
||||
}
|
||||
for _, id := range ids {
|
||||
if id != want {
|
||||
t.Fatalf("provider session ids %v, want %q", ids, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
if recs, err := store.Load(store.MainPath(want)); err != nil || len(recs) == 0 {
|
||||
t.Fatalf("transcript %q not written under the session id: %d records, err=%v", want, len(recs), err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRefusesDamagedTranscript(t *testing.T) {
|
||||
f := newSessionFixture(t)
|
||||
for _, kind := range []string{"worker", "mainagent"} {
|
||||
for _, tc := range damageCases {
|
||||
t.Run(kind+"/"+tc.name, func(t *testing.T) {
|
||||
store := transcript.NewStore(t.TempDir())
|
||||
id := WorkerSessionID(f.expID, f.intentID)
|
||||
if kind == "mainagent" {
|
||||
id = fmt.Sprintf("exp%d-main", f.expID)
|
||||
}
|
||||
original := damageTranscript(t, store, id, tc.unreadable)
|
||||
p := &sessionProbe{}
|
||||
_, err := f.run(t, kind, p, store)
|
||||
assertRestoreError(t, err, len(p.seen()))
|
||||
assertTranscriptUntouched(t, store, id, original, tc.unreadable)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
+10
-2
@@ -540,9 +540,12 @@ func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.
|
||||
if hooks != nil { // typed-nil guard: only set when concrete (avoids harness panic)
|
||||
opts.Hooks = hooks
|
||||
}
|
||||
// The session id is the provider-side identity (session header, usage metering)
|
||||
// with or without persistence; Norma only puts it on ctx when it writes a transcript.
|
||||
opts.SessionID = WorkerSessionID(ts.ID(), intent.ID)
|
||||
ctx = transcript.WithSessionID(ctx, opts.SessionID)
|
||||
if w.tx != nil { // persist raw LLM conversation; one file per worked intent
|
||||
opts.Transcript = w.tx
|
||||
opts.SessionID = WorkerSessionID(ts.ID(), intent.ID)
|
||||
}
|
||||
intentID := intent.ID
|
||||
emitWrap := func(r db.Activity) {
|
||||
@@ -574,7 +577,12 @@ func (w *Worker) execute(ctx context.Context, name string, taskID int64, as *db.
|
||||
// restarting from scratch.
|
||||
alreadyRecorded := false
|
||||
if w.tx != nil {
|
||||
_ = s.Resume(opts.SessionID)
|
||||
// A corrupt or unreadable transcript stops the run before any model or tool
|
||||
// call. No terminal reason: the engine marks the intent blocked without the
|
||||
// model_error reruns, which would only fail the same way again.
|
||||
if err := resumeSession(s, opts.SessionID); err != nil {
|
||||
return "", tsx.Writes(), err
|
||||
}
|
||||
alreadyRecorded = requestID != "" && hasWorkerChatMessage(s.Messages(), requestID)
|
||||
if len(s.Messages()) > 0 && message == "" {
|
||||
seedUnlockFromHistory(s.Messages(), def.UnlockSkill)
|
||||
|
||||
+6
-5
@@ -341,11 +341,12 @@ func applyCostEur(d *DB, out []ProfileUsage) error {
|
||||
if !ok || (pr.in == 0 && pr.outp == 0 && pr.cr == 0 && pr.cw == 0) {
|
||||
continue
|
||||
}
|
||||
// InputTokens includes the cached subset (Anthropic and OpenAI usage
|
||||
// semantics), so the fresh share is input minus cache reads; charging
|
||||
// both buckets at full input price double-counts the cache. Clamped at
|
||||
// zero for gateways that report the buckets disjointly.
|
||||
fresh := out[i].InputTokens - out[i].CacheReadTokens
|
||||
// InputTokens already includes both cached subsets: Norma normalizes it to
|
||||
// input + cache read + cache write (OpenAI prompt_tokens semantics). The
|
||||
// fresh share is therefore input minus cache reads minus cache writes;
|
||||
// charging either subset at full input price as well bills it twice.
|
||||
// Clamped at zero for gateways that report the buckets disjointly.
|
||||
fresh := out[i].InputTokens - out[i].CacheReadTokens - out[i].CacheWriteTokens
|
||||
if fresh < 0 {
|
||||
fresh = 0
|
||||
}
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// costProfiles saves a priced profile (EUR per million tokens: input 2.00,
|
||||
// output 4.00, cache read 0.20, cache write 2.50) and an unpriced one.
|
||||
func costProfiles(t *testing.T, d *DB) (priced, unpriced string) {
|
||||
t.Helper()
|
||||
suffix := time.Now().UnixNano()
|
||||
priced, unpriced = fmt.Sprintf("t-cost-priced-%d", suffix), fmt.Sprintf("t-cost-unpriced-%d", suffix)
|
||||
for _, p := range []*LLMProfile{
|
||||
{Name: priced, Format: "anthropic", Model: "m", APIKey: "k",
|
||||
PriceInputPerM: 2, PriceOutputPerM: 4, PriceCacheReadPerM: 0.2, PriceCacheWritePerM: 2.5},
|
||||
{Name: unpriced, Format: "anthropic", Model: "m", APIKey: "k"},
|
||||
} {
|
||||
id, err := d.SaveProfile(p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
|
||||
}
|
||||
return priced, unpriced
|
||||
}
|
||||
|
||||
func assertCost(t *testing.T, got *float64, want float64) {
|
||||
t.Helper()
|
||||
if got == nil {
|
||||
t.Fatalf("cost=nil, want %.2f", want)
|
||||
}
|
||||
if math.Abs(*got-want) > 1e-9 {
|
||||
t.Fatalf("cost=%.4f, want %.2f", *got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// InputTokens already includes the cache read and cache write subsets, so only
|
||||
// the remainder is billed at the input price.
|
||||
func TestCostSeparatesFreshReadAndWriteBuckets(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
priced, unpriced := costProfiles(t, d)
|
||||
// Every case adds 100k output tokens: 0.1 × 4.00 = 0.40.
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
in, read, write int
|
||||
want float64
|
||||
}{
|
||||
{"fresh", 1_000_000, 0, 0, 2.40}, // 1M×2.00 + 0.40
|
||||
{"read_only", 1_000_000, 1_000_000, 0, 0.60}, // 0 fresh + 1M×0.20 + 0.40
|
||||
{"write_only", 1_000_000, 0, 1_000_000, 2.90}, // 0 fresh + 1M×2.50 + 0.40
|
||||
{"mixed", 1_000_000, 200_000, 300_000, 2.19}, // 500k×2.00 + 200k×0.20 + 300k×2.50 + 0.40
|
||||
{"inconsistent_gateway", 100_000, 200_000, 300_000, 1.19}, // fresh clamped to 0 + 0.04 + 0.75 + 0.40
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
out := []ProfileUsage{
|
||||
{ProfileName: priced, InputTokens: tc.in, CacheReadTokens: tc.read, CacheWriteTokens: tc.write, OutputTokens: 100_000},
|
||||
{ProfileName: unpriced, InputTokens: tc.in, CacheReadTokens: tc.read, CacheWriteTokens: tc.write, OutputTokens: 100_000},
|
||||
{ProfileName: "t-cost-unknown-profile", InputTokens: tc.in, OutputTokens: 100_000},
|
||||
}
|
||||
if err := applyCostEur(d, out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertCost(t, out[0].CostEur, tc.want)
|
||||
if out[1].CostEur != nil || out[2].CostEur != nil {
|
||||
t.Fatalf("unpriced usage got a cost: %v %v", out[1].CostEur, out[2].CostEur)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Several ledger rows of one profile are summed before pricing; with no clamp
|
||||
// involved the profile cost equals the sum of the per-row costs.
|
||||
func TestUsageByProfileCostAggregatesRows(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
priced, unpriced := costProfiles(t, d)
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_usage WHERE profile_name IN ($1,$2)`, priced, unpriced) })
|
||||
rows := []LLMUsage{
|
||||
// fresh 400k×2.00 + 200k×0.20 + 40k×4.00 = 0.80 + 0.04 + 0.16 = 1.00
|
||||
{InputTokens: 600_000, CacheRead: 200_000, OutputTokens: 40_000},
|
||||
// fresh 100k×2.00 + 300k×2.50 + 60k×4.00 = 0.20 + 0.75 + 0.24 = 1.19
|
||||
{InputTokens: 400_000, CacheWrite: 300_000, OutputTokens: 60_000},
|
||||
}
|
||||
for _, profile := range []string{priced, unpriced} {
|
||||
for _, r := range rows {
|
||||
r.ProfileName, r.Model, r.Status = profile, "m", "ok"
|
||||
if err := d.InsertLLMUsage(&r); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
usage, err := d.UsageByProfile()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := 0
|
||||
for _, p := range usage {
|
||||
switch p.ProfileName {
|
||||
case priced:
|
||||
found++
|
||||
if p.Calls != 2 || p.InputTokens != 1_000_000 || p.CacheReadTokens != 200_000 || p.CacheWriteTokens != 300_000 || p.OutputTokens != 100_000 {
|
||||
t.Fatalf("aggregate %+v", p)
|
||||
}
|
||||
assertCost(t, p.CostEur, 2.19) // = 1.00 + 1.19, the mixed case
|
||||
case unpriced:
|
||||
found++
|
||||
if p.Calls != 2 || p.CostEur != nil {
|
||||
t.Fatalf("unpriced aggregate %+v", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
if found != 2 {
|
||||
t.Fatalf("profiles missing from UsageByProfile: found %d of 2", found)
|
||||
}
|
||||
}
|
||||
+6
-6
@@ -195,11 +195,11 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
|
||||
return
|
||||
}
|
||||
// Always track usage (cheap); accumulate text/thinking only for bodies.
|
||||
// Anthropic (and the other providers) split token usage across events:
|
||||
// message_start carries ONLY the input side (input + cache), message_delta
|
||||
// ONLY the output. They must be FOLDED with Add — overwriting on delta
|
||||
// would zero out the input already counted at start (mirrors the SDK's own
|
||||
// llm.Accumulator; see norma/llm/accumulate.go).
|
||||
// message_start carries the input side (input + cache) and is added.
|
||||
// Since Norma v0.4.7 message_delta also echoes that input side next to the
|
||||
// output, so it is folded with llm.Usage.MergeEcho (max on echoed fields,
|
||||
// sum on output), as llm.Accumulator does; Add would count the input and
|
||||
// cache twice. An output-only delta folds the same either way.
|
||||
switch ev.Type {
|
||||
case llm.SETextDelta:
|
||||
if recordBodies {
|
||||
@@ -215,7 +215,7 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
|
||||
if ev.StopReason != "" {
|
||||
stopReason = ev.StopReason
|
||||
}
|
||||
usage.Add(ev.Usage)
|
||||
usage.MergeEcho(ev.Usage)
|
||||
}
|
||||
if !yield(ev, err) {
|
||||
return
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
package llmrec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// echoStreamProvider streams one message_start and one message_delta with the
|
||||
// given (already normalized) usage, as Norma's providers emit them.
|
||||
type echoStreamProvider struct {
|
||||
completeProvider
|
||||
start, delta llm.Usage
|
||||
}
|
||||
|
||||
func (p echoStreamProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: p.start},
|
||||
{Type: llm.SETextDelta, Text: "answer"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: p.delta},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Norma normalizes Anthropic usage so InputTokens includes the cache read and
|
||||
// write subsets, and since v0.4.7 message_delta echoes that input side.
|
||||
var usageEchoCases = []struct {
|
||||
name string
|
||||
start, delta llm.Usage
|
||||
want llm.Usage
|
||||
}{
|
||||
{
|
||||
name: "full_echo",
|
||||
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
delta: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
|
||||
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
},
|
||||
{
|
||||
// Raw start: input 500, cache read 200, cache creation 40 → normalized 740.
|
||||
// Raw delta: input 500, output 77, no cache fields → normalized 500.
|
||||
name: "partial_echo",
|
||||
start: llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
delta: llm.Usage{InputTokens: 500, OutputTokens: 77},
|
||||
want: llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
},
|
||||
{
|
||||
name: "no_echo",
|
||||
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
delta: llm.Usage{OutputTokens: 5},
|
||||
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
},
|
||||
}
|
||||
|
||||
func drainRecorderStream(t *testing.T, ctx context.Context, r *Recorder) {
|
||||
t.Helper()
|
||||
for _, err := range r.Stream(ctx, llm.CompletionRequest{}) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The metered usage of a streamed call counts an echoed input side once.
|
||||
func TestStreamUsageEchoCountedOnce(t *testing.T) {
|
||||
for _, tc := range usageEchoCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var got []llm.Usage
|
||||
ctx := WithUsageSink(t.Context(), func(u llm.Usage) { got = append(got, u) })
|
||||
drainRecorderStream(t, ctx, Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, nil, "fixture", "profile", "", "", func() bool { return false }))
|
||||
if len(got) != 1 || got[0] != tc.want {
|
||||
t.Fatalf("metered %+v, want one call with %+v", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The llm_usage ledger gets one row per streamed call, with the echo counted once.
|
||||
func TestStreamUsageEchoLedgerRow(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v)", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
if err := pg.EnsureLLMUsageTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tc := range usageEchoCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
profile := "usage-echo-" + uuid.NewString()
|
||||
t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE profile_name=$1`, profile) })
|
||||
drainRecorderStream(t, t.Context(), Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, pg, "fixture", profile, "", "", func() bool { return false }))
|
||||
var rows int
|
||||
var got llm.Usage
|
||||
err := pg.QueryRow(`SELECT count(*), COALESCE(sum(input_tokens),0), COALESCE(sum(output_tokens),0), COALESCE(sum(cache_read),0), COALESCE(sum(cache_write),0)
|
||||
FROM llm_usage WHERE profile_name=$1`, profile).Scan(&rows, &got.InputTokens, &got.OutputTokens, &got.CacheReadTokens, &got.CacheWriteTokens)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rows != 1 || got != tc.want {
|
||||
t.Fatalf("ledger rows=%d usage=%+v, want 1 row with %+v", rows, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+75
-19
@@ -317,10 +317,14 @@ func (r *taskLLMRuntime) Complete(ctx context.Context, req llm.CompletionRequest
|
||||
}
|
||||
|
||||
func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) (llm.Message, string, llm.Usage, error) {
|
||||
// total is the usage of the logical call: the sum of every physical attempt
|
||||
// (same-provider retries and profile failover). Each attempt is already billed
|
||||
// once by the recorder below; this sum only feeds the caller's accounting.
|
||||
var total llm.Usage
|
||||
for {
|
||||
selection, err := hooks.current(ctx)
|
||||
if err != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, err
|
||||
return llm.Message{}, "", total, err
|
||||
}
|
||||
var (
|
||||
msg llm.Message
|
||||
@@ -333,13 +337,14 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
|
||||
retries, backoffOf := sameProviderRetryPolicy(selection.retry)
|
||||
for attempt := 0; ; attempt++ {
|
||||
msg, sr, usage, callErr = selection.provider.Complete(ctx, req)
|
||||
total.Add(usage)
|
||||
if callErr != nil && ctx.Err() == nil &&
|
||||
attempt < retries && isRetryableStreamError(callErr) {
|
||||
backoff := backoffOf(attempt)
|
||||
log.Printf("[task-llm] task %s non-streaming call failed, retrying on the same provider after %v (%d/%d): %v",
|
||||
taskID, backoff, attempt+1, retries, callErr)
|
||||
if sleepCtx(ctx, backoff) {
|
||||
break // ctx cancelled during backoff → stop retrying
|
||||
return llm.Message{}, "", total, ctx.Err() // cancelled during backoff → stop
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -347,7 +352,7 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
|
||||
}
|
||||
if callErr == nil {
|
||||
resetLLMFailStreak(taskID, selection.profileID)
|
||||
return msg, sr, usage, nil
|
||||
return msg, sr, total, nil
|
||||
}
|
||||
// Non-quota failures on a chain profile escalate after a consecutive
|
||||
// streak (#99): same advance path as quota exhaust, truthful reason.
|
||||
@@ -356,31 +361,31 @@ func completeTaskLLM(ctx context.Context, taskID string, req llm.CompletionReque
|
||||
cause := fmt.Errorf("profile #%d failed %d consecutive calls (last: %w)", selection.profileID, llmFailStreakThreshold, callErr)
|
||||
transition, markErr := hooks.exhaust(selection, cause)
|
||||
if markErr != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile exhausted after %v: %w", callErr, markErr)
|
||||
return llm.Message{}, "", total, fmt.Errorf("mark profile exhausted after %v: %w", callErr, markErr)
|
||||
}
|
||||
if transition.Advanced && !transition.Stale && hooks.transition != nil {
|
||||
hooks.transition(selection, transition, callErr)
|
||||
}
|
||||
if !transition.Stale && transition.NextProfileID == nil {
|
||||
return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
return llm.Message{}, "", total, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
}
|
||||
continue // replay the same logical request on the next profile
|
||||
}
|
||||
return llm.Message{}, "", llm.Usage{}, callErr
|
||||
return llm.Message{}, "", total, callErr
|
||||
}
|
||||
// profileID=0: the explicit chain was cleared; non-quota error: pass through. Neither changes the task failover state.
|
||||
if selection.profileID == 0 || !isQuotaExhaustedError(callErr) {
|
||||
return llm.Message{}, "", llm.Usage{}, callErr
|
||||
return llm.Message{}, "", total, callErr
|
||||
}
|
||||
transition, markErr := hooks.exhaust(selection, callErr)
|
||||
if markErr != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr)
|
||||
return llm.Message{}, "", total, fmt.Errorf("mark profile quota exhausted after %v: %w", callErr, markErr)
|
||||
}
|
||||
if transition.Advanced && !transition.Stale && hooks.transition != nil {
|
||||
hooks.transition(selection, transition, callErr)
|
||||
}
|
||||
if !transition.Stale && transition.NextProfileID == nil {
|
||||
return llm.Message{}, "", llm.Usage{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
return llm.Message{}, "", total, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: callErr}
|
||||
}
|
||||
// No output has been delivered to the caller, so replaying the same logical request on the next profile is safe.
|
||||
}
|
||||
@@ -397,6 +402,35 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
committed := false
|
||||
var pending []llm.StreamEvent
|
||||
var streamErr error
|
||||
// The caller's llm.Accumulator sums usage across the whole logical call
|
||||
// (every attempt and profile): message_start is added, message_delta is
|
||||
// folded with MergeEcho. Folding one attempt's echo against the sum of
|
||||
// earlier attempts would lose input/cache tokens, so each attempt's echoes
|
||||
// are normalized here against that attempt alone and only the increment
|
||||
// is passed on, as a message_start. The recorder below already billed the
|
||||
// physical attempt; this only shapes what the caller adds up.
|
||||
var attemptUsage llm.Usage
|
||||
forward := func(event llm.StreamEvent) bool {
|
||||
switch event.Type {
|
||||
case llm.SEMessageStart:
|
||||
attemptUsage.Add(event.Usage)
|
||||
case llm.SEMessageDelta:
|
||||
next := attemptUsage
|
||||
next.MergeEcho(event.Usage)
|
||||
inc := llm.Usage{
|
||||
InputTokens: next.InputTokens - attemptUsage.InputTokens,
|
||||
OutputTokens: next.OutputTokens - attemptUsage.OutputTokens,
|
||||
CacheReadTokens: next.CacheReadTokens - attemptUsage.CacheReadTokens,
|
||||
CacheWriteTokens: next.CacheWriteTokens - attemptUsage.CacheWriteTokens,
|
||||
}
|
||||
attemptUsage = next
|
||||
if inc != (llm.Usage{}) && !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: inc}, nil) {
|
||||
return false
|
||||
}
|
||||
event.Usage = llm.Usage{}
|
||||
}
|
||||
return yield(event, nil)
|
||||
}
|
||||
retries, backoffOf := sameProviderRetryPolicy(selection.retry)
|
||||
// Same-provider safe-window retry: a transient failure before committed
|
||||
// (nothing delivered to the caller yet) can be replayed verbatim without
|
||||
@@ -405,6 +439,7 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
committed = false
|
||||
pending = nil
|
||||
streamErr = nil
|
||||
attemptUsage = llm.Usage{}
|
||||
for event, err := range selection.provider.Stream(ctx, req) {
|
||||
if err != nil {
|
||||
streamErr = err
|
||||
@@ -416,24 +451,34 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
}
|
||||
if !committed {
|
||||
for _, buffered := range pending {
|
||||
if !yield(buffered, nil) {
|
||||
if !forward(buffered) {
|
||||
return
|
||||
}
|
||||
}
|
||||
pending = nil
|
||||
committed = true
|
||||
}
|
||||
if !yield(event, nil) {
|
||||
if !forward(event) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if streamErr != nil && !committed {
|
||||
// A failed uncommitted attempt delivers none of its events (they are
|
||||
// never replayed), but the usage it reported survives as one
|
||||
// message_start so the logical call still counts it.
|
||||
if u := pendingUsage(pending); u != (llm.Usage{}) && !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: u}, nil) {
|
||||
return
|
||||
}
|
||||
pending = nil
|
||||
}
|
||||
if streamErr != nil && !committed && ctx.Err() == nil &&
|
||||
attempt < retries && isRetryableStreamError(streamErr) {
|
||||
backoff := backoffOf(attempt)
|
||||
log.Printf("[task-llm] task %s pre-commit stream failure, retrying on the same provider after %v (%d/%d): %v",
|
||||
taskID, backoff, attempt+1, retries, streamErr)
|
||||
if sleepCtx(ctx, backoff) {
|
||||
break // ctx cancelled during backoff → stop retrying
|
||||
yield(llm.StreamEvent{}, ctx.Err()) // cancelled during backoff → stop
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -442,7 +487,7 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
if streamErr == nil {
|
||||
resetLLMFailStreak(taskID, selection.profileID)
|
||||
for _, buffered := range pending {
|
||||
if !yield(buffered, nil) {
|
||||
if !forward(buffered) {
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -474,11 +519,6 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
// bundle was still in use. Agent/global fallback errors follow the legacy
|
||||
// behavior and never mutate task failover state.
|
||||
if selection.profileID == 0 || !isQuotaExhaustedError(streamErr) {
|
||||
for _, buffered := range pending {
|
||||
if !yield(buffered, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
yield(llm.StreamEvent{}, streamErr)
|
||||
return
|
||||
}
|
||||
@@ -508,9 +548,25 @@ func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest
|
||||
}
|
||||
}
|
||||
|
||||
// pendingUsage is the usage the buffered (uncommitted) events of one attempt
|
||||
// reported. Only message_start can be buffered: message_delta commits.
|
||||
func pendingUsage(pending []llm.StreamEvent) llm.Usage {
|
||||
var u llm.Usage
|
||||
for _, ev := range pending {
|
||||
if ev.Type == llm.SEMessageStart {
|
||||
u.Add(ev.Usage)
|
||||
}
|
||||
}
|
||||
return u
|
||||
}
|
||||
|
||||
// streamEventCommitsOutput reports whether delivering event makes the attempt
|
||||
// unreplayable. Signed and redacted thinking and tool-use signatures are opaque
|
||||
// output the model bound to this exact response, so they commit like text.
|
||||
func streamEventCommitsOutput(event llm.StreamEvent) bool {
|
||||
switch event.Type {
|
||||
case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON:
|
||||
case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON,
|
||||
llm.SEThinkingSignature, llm.SERedactedThinking, llm.SEToolUseSignature:
|
||||
return event.Text != ""
|
||||
case llm.SEToolUseStart, llm.SEMessageDelta, llm.SEMessageStop:
|
||||
return true
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/llmrec"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// Usage reported by one failed attempt, and by a successful one. A stream
|
||||
// reports the input side in message_start and echoes it, with the output, in
|
||||
// message_delta; the atomic path returns the attempt's normalized usage.
|
||||
var (
|
||||
failedAttemptUsage = llm.Usage{InputTokens: 11, OutputTokens: 2, CacheReadTokens: 3, CacheWriteTokens: 2}
|
||||
okStartUsage = llm.Usage{InputTokens: 13, CacheReadTokens: 3, CacheWriteTokens: 2}
|
||||
okDeltaUsage = llm.Usage{InputTokens: 17, OutputTokens: 5, CacheReadTokens: 4, CacheWriteTokens: 2}
|
||||
okAttemptUsage = llm.Usage{InputTokens: 17, OutputTokens: 5, CacheReadTokens: 4, CacheWriteTokens: 2}
|
||||
)
|
||||
|
||||
// usageAttemptProvider fails its first fail attempts (always, when fail < 0)
|
||||
// after reporting failedAttemptUsage, then answers "done".
|
||||
type usageAttemptProvider struct {
|
||||
fail int
|
||||
failErr error
|
||||
onFail func()
|
||||
// extra events a failed stream attempt emits before its error.
|
||||
failEvents []llm.StreamEvent
|
||||
calls int
|
||||
}
|
||||
|
||||
func (p *usageAttemptProvider) failing() bool {
|
||||
p.calls++
|
||||
failed := p.fail < 0 || p.calls <= p.fail
|
||||
if failed && p.onFail != nil {
|
||||
p.onFail()
|
||||
}
|
||||
return failed
|
||||
}
|
||||
|
||||
func (p *usageAttemptProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
if p.failing() {
|
||||
return llm.Message{}, "", failedAttemptUsage, p.failErr
|
||||
}
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("done")}}, "end_turn", okAttemptUsage, nil
|
||||
}
|
||||
|
||||
func (p *usageAttemptProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
if p.failing() {
|
||||
events := append([]llm.StreamEvent{{Type: llm.SEMessageStart, Usage: failedAttemptUsage}}, p.failEvents...)
|
||||
for _, ev := range events {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
yield(llm.StreamEvent{}, p.failErr)
|
||||
return
|
||||
}
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: okStartUsage},
|
||||
{Type: llm.SETextDelta, Text: "done"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: okDeltaUsage},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// logicalResult is what the SDK sees for one logical call.
|
||||
type logicalResult struct {
|
||||
usage llm.Usage
|
||||
text string
|
||||
events []llm.StreamEvent
|
||||
err error
|
||||
}
|
||||
|
||||
func runLogicalCall(ctx context.Context, stream bool, taskID string, hooks taskLLMStreamHooks) logicalResult {
|
||||
if !stream {
|
||||
msg, _, usage, err := completeTaskLLM(ctx, taskID, llm.CompletionRequest{}, hooks)
|
||||
return logicalResult{usage: usage, text: msg.Text(), err: err}
|
||||
}
|
||||
var r logicalResult
|
||||
acc := llm.NewAccumulator()
|
||||
for ev, err := range streamTaskLLM(ctx, taskID, llm.CompletionRequest{}, hooks) {
|
||||
if err != nil {
|
||||
r.err = err
|
||||
break
|
||||
}
|
||||
r.events = append(r.events, ev)
|
||||
acc.Add(ev)
|
||||
}
|
||||
r.usage, r.text = acc.Usage, acc.Message().Text()
|
||||
return r
|
||||
}
|
||||
|
||||
func singleProfileHooks(p llm.Provider, retry agent.RetryConfig) taskLLMStreamHooks {
|
||||
return taskLLMStreamHooks{
|
||||
current: func(context.Context) (taskLLMSelection, error) {
|
||||
return taskLLMSelection{provider: p, retry: retry}, nil
|
||||
},
|
||||
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
|
||||
return db.TaskLLMTransition{}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func sumUsage(us ...llm.Usage) llm.Usage {
|
||||
var t llm.Usage
|
||||
for _, u := range us {
|
||||
t.Add(u)
|
||||
}
|
||||
return t
|
||||
}
|
||||
|
||||
var usageModes = []struct {
|
||||
name string
|
||||
stream bool
|
||||
}{{"atomic", false}, {"stream", true}}
|
||||
|
||||
func TestTaskLLMRetrySumsAttemptUsage(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
p := &usageAttemptProvider{fail: 1, failErr: errors.New("connection reset by peer")}
|
||||
r := runLogicalCall(t.Context(), m.stream, "p12-retry", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
want := sumUsage(failedAttemptUsage, okAttemptUsage) // {28 7 7 4}
|
||||
if r.usage != want || p.calls != 2 || r.text != "done" {
|
||||
t.Fatalf("usage=%+v want %+v, calls=%d, text=%q", r.usage, want, p.calls, r.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMExhaustedRetriesReturnSummedUsage(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
p := &usageAttemptProvider{fail: -1, failErr: errors.New("connection reset by peer")}
|
||||
r := runLogicalCall(t.Context(), m.stream, "p12-exhausted", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err == nil || r.err.Error() != "connection reset by peer" {
|
||||
t.Fatalf("err=%v", r.err)
|
||||
}
|
||||
want := sumUsage(failedAttemptUsage, failedAttemptUsage)
|
||||
if r.usage != want || p.calls != 2 {
|
||||
t.Fatalf("usage=%+v want %+v, calls=%d", r.usage, want, p.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMProfileFailoverSumsUsage(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
first := &usageAttemptProvider{fail: -1, failErr: errors.New("openai: status 402: insufficient_quota")}
|
||||
second := &usageAttemptProvider{}
|
||||
active := int64(11)
|
||||
hooks := taskLLMStreamHooks{
|
||||
current: func(context.Context) (taskLLMSelection, error) {
|
||||
if active == 11 {
|
||||
return taskLLMSelection{profileID: 11, provider: first}, nil
|
||||
}
|
||||
return taskLLMSelection{profileID: 22, provider: second}, nil
|
||||
},
|
||||
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
|
||||
active = 22
|
||||
next := active
|
||||
return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil
|
||||
},
|
||||
}
|
||||
r := runLogicalCall(t.Context(), m.stream, "p12-failover-"+m.name, hooks)
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
want := sumUsage(failedAttemptUsage, okAttemptUsage)
|
||||
if r.usage != want || first.calls != 1 || second.calls != 1 || r.text != "done" {
|
||||
t.Fatalf("usage=%+v want %+v, calls first=%d second=%d, text=%q", r.usage, want, first.calls, second.calls, r.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMCancelDuringBackoffStops(t *testing.T) {
|
||||
for _, m := range usageModes {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(t.Context())
|
||||
defer cancel()
|
||||
p := &usageAttemptProvider{fail: -1, failErr: errors.New("connection reset by peer")}
|
||||
// Cancel once the retry is waiting out its (long) backoff.
|
||||
p.onFail = func() { time.AfterFunc(20*time.Millisecond, cancel) }
|
||||
r := runLogicalCall(ctx, m.stream, "p12-cancel", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 2, StreamInterval: time.Minute}))
|
||||
if !errors.Is(r.err, context.Canceled) {
|
||||
t.Fatalf("err=%v, want context.Canceled", r.err)
|
||||
}
|
||||
if p.calls != 1 || r.usage != failedAttemptUsage {
|
||||
t.Fatalf("calls=%d usage=%+v", p.calls, r.usage)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMUncommittedStreamAttemptKeepsUsageWithoutReplay(t *testing.T) {
|
||||
p := &usageAttemptProvider{
|
||||
fail: 1,
|
||||
failErr: errors.New("anthropic: overloaded_error"),
|
||||
// Buffered, non-committing events of the failed attempt.
|
||||
failEvents: []llm.StreamEvent{{Type: llm.SEThinkingDelta}, {Type: llm.SEToolInputJSON}},
|
||||
}
|
||||
r := runLogicalCall(t.Context(), true, "p12-uncommitted", singleProfileHooks(p, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
if want := sumUsage(failedAttemptUsage, okAttemptUsage); r.usage != want {
|
||||
t.Fatalf("usage=%+v want %+v", r.usage, want)
|
||||
}
|
||||
if r.text != "done" {
|
||||
t.Fatalf("text=%q, want the second attempt's only", r.text)
|
||||
}
|
||||
// The failed attempt contributes one usage-only message_start, nothing else.
|
||||
if len(r.events) == 0 || r.events[0].Type != llm.SEMessageStart || r.events[0].Usage != failedAttemptUsage {
|
||||
t.Fatalf("first event %+v", r.events)
|
||||
}
|
||||
for _, ev := range r.events {
|
||||
if ev.Type == llm.SEThinkingDelta || ev.Type == llm.SEToolInputJSON {
|
||||
t.Fatalf("buffered event of the failed attempt was delivered: %+v", r.events)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskLLMOpaqueThinkingCommitsStream(t *testing.T) {
|
||||
defer withZeroRetryBackoff()()
|
||||
for _, kind := range []llm.StreamEventType{llm.SEThinkingSignature, llm.SERedactedThinking, llm.SEToolUseSignature} {
|
||||
t.Run(string(kind), func(t *testing.T) {
|
||||
if !streamEventCommitsOutput(llm.StreamEvent{Type: kind, Text: "opaque"}) {
|
||||
t.Fatalf("%s does not commit", kind)
|
||||
}
|
||||
// A retryable failure after the opaque block is surfaced, not replayed.
|
||||
p := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: kind, Text: "opaque"}}, err: errors.New("connection reset by peer")}
|
||||
events, err := collectTaskLLMStream(streamTaskLLM(t.Context(), "p12-opaque", llm.CompletionRequest{}, singleProfileHooks(p, agent.RetryConfig{})))
|
||||
if err == nil || p.calls != 1 || len(events) != 1 || events[0].Type != kind {
|
||||
t.Fatalf("calls=%d events=%+v err=%v", p.calls, events, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The recorder below the task router writes one llm_usage row per physical
|
||||
// attempt; summing the attempts for the logical call adds no row.
|
||||
func TestTaskLLMRetryKeepsOneLedgerRowPerPhysicalCall(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) — skipping", err)
|
||||
}
|
||||
defer m.Close()
|
||||
pg := m.PG()
|
||||
for _, mode := range usageModes {
|
||||
t.Run(mode.name, func(t *testing.T) {
|
||||
taskID := fmt.Sprintf("p12-ledger-%s-%d", mode.name, time.Now().UnixNano())
|
||||
t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE task_id=$1`, taskID) })
|
||||
p := &usageAttemptProvider{fail: 1, failErr: errors.New("connection reset by peer")}
|
||||
rec := llmrec.Wrap(p, pg, "test-model", "test-profile", "", "", func() bool { return false })
|
||||
ctx := llmrec.WithTaskID(t.Context(), taskID)
|
||||
r := runLogicalCall(ctx, mode.stream, taskID, singleProfileHooks(rec, agent.RetryConfig{StreamAttempts: 1, StreamInterval: time.Millisecond}))
|
||||
if r.err != nil {
|
||||
t.Fatal(r.err)
|
||||
}
|
||||
var rows, failed int
|
||||
if err := pg.QueryRow(`SELECT count(*), count(*) FILTER (WHERE status='error') FROM llm_usage WHERE task_id=$1`, taskID).Scan(&rows, &failed); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rows != 2 || failed != 1 || p.calls != 2 {
|
||||
t.Fatalf("llm_usage rows=%d (error %d), physical calls=%d; want 2 rows, 1 error", rows, failed, p.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -32,8 +32,11 @@ func (s SideQuestionService) Answer(ctx context.Context, req llm.CompletionReque
|
||||
out.Text += ev.Text
|
||||
case llm.SEToolUseStart:
|
||||
out.ToolUse = true
|
||||
case llm.SEMessageStart, llm.SEMessageDelta:
|
||||
case llm.SEMessageStart:
|
||||
out.Usage.Add(ev.Usage)
|
||||
case llm.SEMessageDelta:
|
||||
// message_delta may echo message_start's input side: MergeEcho, not Add.
|
||||
out.Usage.MergeEcho(ev.Usage)
|
||||
case llm.SEMessageStop:
|
||||
complete = true
|
||||
}
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
package sidequestion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// A streamed side answer counts message_delta's echo of the input side once.
|
||||
func TestAnswerStreamUsageEchoCountedOnce(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
start, delta llm.Usage
|
||||
want llm.Usage
|
||||
}{
|
||||
{"full_echo",
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
|
||||
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
|
||||
{"partial_echo",
|
||||
llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
llm.Usage{InputTokens: 500, OutputTokens: 77},
|
||||
llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40}},
|
||||
{"no_echo",
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
llm.Usage{OutputTokens: 5},
|
||||
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
p := fakeProvider{stream: func(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: tc.start},
|
||||
{Type: llm.SETextDelta, Text: "answer"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: tc.delta},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}}
|
||||
out, err := (SideQuestionService{Provider: p}).Answer(t.Context(), llm.CompletionRequest{}, true, nil)
|
||||
if err != nil || out.Text != "answer" || out.Usage != tc.want {
|
||||
t.Fatalf("answer=%+v err=%v, want usage %+v", out, err, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Generated
+588
-234
File diff suppressed because it is too large
Load Diff
+2
-3
@@ -21,7 +21,6 @@
|
||||
},
|
||||
"dependencies": {
|
||||
"@antv/g6": "^5.1.1",
|
||||
"@antv/graphin": "^3.0.5",
|
||||
"@base-ui/react": "^1.6.0",
|
||||
"@dnd-kit/core": "^6.3.1",
|
||||
"@dnd-kit/modifiers": "^9.0.0",
|
||||
@@ -40,7 +39,7 @@
|
||||
"geist": "^1.7.2",
|
||||
"input-otp": "^1.4.2",
|
||||
"lucide-react": "^1.22.0",
|
||||
"next": "^16.2.9",
|
||||
"next": "^16.4.6",
|
||||
"next-themes": "^0.4.6",
|
||||
"radix-ui": "^1.6.0",
|
||||
"react": "^19.2.7",
|
||||
@@ -52,7 +51,6 @@
|
||||
"react-resizable-panels": "^4.12.0",
|
||||
"recharts": "^3.8.0",
|
||||
"remark-gfm": "^4.0.1",
|
||||
"shadcn": "^4.12.0",
|
||||
"simple-icons": "^16.24.1",
|
||||
"sonner": "^2.0.7",
|
||||
"tailwind-merge": "^3.6.0",
|
||||
@@ -74,6 +72,7 @@
|
||||
"husky": "^9.1.7",
|
||||
"lint-staged": "^16.4.0",
|
||||
"postcss": "^8.5.16",
|
||||
"shadcn": "^4.12.0",
|
||||
"tailwindcss": "^4.1.5",
|
||||
"ts-node": "^10.9.2",
|
||||
"tw-animate-css": "^1.4.0",
|
||||
|
||||
Reference in New Issue
Block a user