okami.mom/albedo
git clone https://git.nkp.pet/okami.mom/albedo
5b3f8be
main
1//go:build unix 2 3// Prompt scenarios: `albedo -p` sends one prompt, blocks until its turn ends, 4// and prints only the final assistant text, so a script can pipe it. 5package e2e 6 7import ( 8"context" 9"encoding/json" 10"errors" 11"strings" 12"testing" 13"time" 14 15"albedo/cli/internal/config" 16"albedo/cli/internal/daemon" 17) 18 19// A new prompt runs in a fresh session on the model named provider/model, and 20// stdout is exactly the reply; `albedo models` lists that same spelling. 21func TestPromptPrintsTheFinalReplyFromTheNamedModel (t * testing.T ) { 22profile := providerRoute (t ,echoReply ) 23model := profile + "/fixture-model" 24if listed := cli (t ,"models" );! strings .Contains (listed ,model + "\n" ) { 25t .Fatalf ("albedo models does not list %s:\n%s" ,model ,listed ) 26 } 27 28stdout := cli (t ,"-p" ,"print me" ,"--model" ,model ) 29if stdout != "echo: print me\n" { 30t .Fatalf ("albedo -p printed %q" ,stdout ) 31 } 32if got := len (suite .provider .requests (profile ));got != 1 { 33t .Fatalf ("provider %s saw %d requests, want 1" ,profile ,got ) 34 } 35} 36 37// --model on an existing session moves only that session: the turn reaches 38// the other provider, and new sessions still default to the active one. 39func TestPromptToASessionSwitchesItsModelWithoutChangingTheDefault (t * testing.T ) { 40other := "PromptSwitchTarget" 41if err := suite .provider .addProfile (other ,echoReply );err != nil { 42t .Fatal (err ) 43 } 44profile := providerRoute (t ,echoReply ) 45id := newSession (t ,t .TempDir ()) 46 47stdout := cli (t ,"--prompt" ,"switch" ,"--session" ,id ,"--model" ,other + "/fixture-model" ) 48if stdout != "echo: switch\n" { 49t .Fatalf ("albedo -p -s printed %q" ,stdout ) 50 } 51if len (suite .provider .requests (other ))!= 1 || len (suite .provider .requests (profile ))!= 0 { 52t .Fatalf ("turn went to the wrong provider: %s=%d %s=%d" ,other , 53len (suite .provider .requests (other )),profile ,len (suite .provider .requests (profile ))) 54 } 55if session := daemonSession (t ,id );session .Provider != other { 56t .Fatalf ("session provider is %q, want %q" ,session .Provider ,other ) 57 } 58profiles ,err := daemon .ProviderProfiles (context .Background (),conn (t )) 59if err != nil || profiles .Active != profile { 60t .Fatalf ("default provider became %q (%v), want %q" ,profiles .Active ,err ,profile ) 61 } 62} 63 64// A turn that fails must end the command with the failure, not hang or print 65// an empty reply; a model no provider offers fails before any session exists. 66func TestPromptReportsAFailedTurnAndAnUnknownModel (t * testing.T ) { 67profile := t .Name () 68if err := saveAndSelectProvider (context .Background (),conn (t ),profile , config.Settings { 69Extension :"openai" , 70BaseURL :suite .provider .server .URL + "/t/missing" , 71APIKey :"fixture-key" , 72Model :"fixture-model" , 73Protocol :"chat_completions" , 74 });err != nil { 75t .Fatal (err ) 76 } 77 78stdout ,stderr ,err := runCLI ("-p" ,"fail" ,"--model" ,profile + "/fixture-model" ) 79if err == nil || stdout != "" || ! strings .Contains (stderr ,"404" ) { 80t .Fatalf ("failed turn: err=%v stdout=%q stderr=%q" ,err ,stdout ,stderr ) 81 } 82 83before := len (daemonSessions (t )) 84_ ,stderr ,err = runCLI ("-p" ,"never sent" ,"--model" ,"no-such-model" ) 85if err == nil || ! strings .Contains (stderr ,"albedo models" ) { 86t .Fatalf ("unknown model: err=%v stderr=%q" ,err ,stderr ) 87 } 88if after := len (daemonSessions (t ));after != before { 89t .Fatalf ("an unknown model still created a session (%d -> %d)" ,before ,after ) 90 } 91} 92 93// --timeout fails the command once it passes and stops the turn it started: 94// the session is idle long before the stalled provider would have answered. 95func TestPromptTimeoutStopsTheTurn (t * testing.T ) { 96entered ,release := make (chan struct {}),make (chan struct {}) 97t .Cleanup (func () {close (release ) }) 98profile := providerRoute (t ,func (request map [string ]any )string { 99if lastUserText (request )== "stall" { 100close (entered ) 101select { 102case <- release : 103case <- time .After (10 * time .Second ): 104 } 105 } 106return echoReply (request ) 107 }) 108// The deadline tests turn cancellation, not session creation or kernel boot. 109session := newSession (t ,t .TempDir ()) 110cli (t ,"-p" ,"warm the timeout session" ,"--session" ,session ) 111waitIdle (t ,session ,profile ,1 ) 112 113started := time .Now () 114stdout ,stderr ,err := runCLI ("-p" ,"stall" ,"--session" ,session ,"--timeout" ,"1s" ) 115if err == nil || stdout != "" || ! strings .Contains (stderr ,"timed out after 1s" ) { 116t .Fatalf ("timed-out prompt: err=%v stdout=%q stderr=%q" ,err ,stdout ,stderr ) 117 } 118select { 119case <- entered : 120default : 121t .Fatal ("the deadline expired before the stalled turn reached the provider" ) 122 } 123_ ,id ,_ := strings .Cut (strings .TrimSpace (stderr ),"the session is " ) 124if id != session { 125t .Fatalf ("timeout reported session %q, want %q" ,id ,session ) 126 } 127for { 128status ,err := daemon .NewChatClient (conn (t ),id ).GetStatus (context .Background ()) 129if err != nil { 130t .Fatalf ("status of %q: %v" ,id ,err ) 131 } 132if status .Idle { 133break 134 } 135if time .Since (started )> 6 * time .Second { 136t .Fatal ("the timed-out turn is still running" ) 137 } 138time .Sleep (100 * time .Millisecond ) 139 } 140} 141 142// Cancellation while queued must leave the already running turn alone. 143func TestPromptTimeoutWhileQueuedDoesNotInterruptAnotherTurn (t * testing.T ) { 144entered ,release := make (chan struct {}),make (chan struct {}) 145releaseClosed := false 146defer func () { 147if ! releaseClosed { 148close (release ) 149 } 150 }() 151profile := providerRoute (t ,func (request map [string ]any )string { 152if lastUserText (request )== "other turn" { 153close (entered ) 154<- release 155 } 156return echoReply (request ) 157 }) 158t .Parallel () 159id := newSession (t ,t .TempDir ()) 160client := daemon .NewChatClient (conn (t ),id ) 161if _ ,err := client .Send (t .Context (),"other turn" ,nil );err != nil { 162t .Fatal (err ) 163 } 164select { 165case <- entered : 166case <- time .After (10 * time .Second ): 167t .Fatal ("other turn never started" ) 168 } 169stdout ,stderr ,err := runCLI ("--prompt" ,"queued turn" ,"--session" ,id ,"--timeout" ,"1s" ) 170if err == nil || stdout != "" || ! strings .Contains (stderr ,"timed out after 1s" ) { 171t .Fatalf ("queued timeout: %v %q %q" ,err ,stdout ,stderr ) 172 } 173status ,statusErr := client .GetStatus (t .Context ()) 174if statusErr != nil || status .Idle { 175t .Fatalf ("queued cancellation stopped the other turn: %+v %v" ,status ,statusErr ) 176 } 177if ! strings .Contains (stderr ,"queued submission cancelled" )|| strings .Contains (stderr ,"stopped" ) { 178t .Fatalf ("incorrect cancellation report: %s" ,stderr ) 179 } 180_ ,identity ,_ := strings .Cut (stderr ,"operation " ) 181inputID ,_ ,_ := strings .Cut (identity ," can be queried" ) 182receipt ,receiptErr := daemon .GetInput (t .Context (),conn (t ),id ,inputID ) 183if receiptErr != nil || receipt .Admission != "accepted" || receipt .Delivery == nil || * receipt .Delivery != "cancelled" || receipt .Turn != nil { 184t .Fatalf ("timed call did not cancel an admitted queued input: %+v, %v" ,receipt ,receiptErr ) 185 } 186close (release ) 187releaseClosed = true 188waitIdle (t ,id ,profile ,1 ) 189if got := len (suite .provider .requests (profile ));got != 1 { 190t .Fatalf ("cancelled queued prompt still ran: %d requests" ,got ) 191 } 192missing ,err := client .PrepareTurn ("never submitted" ,nil ,nil ,false ) 193if err != nil { 194t .Fatal (err ) 195 } 196_ ,err = client .CancelSubmission (t .Context (),missing .ID ()) 197if failure ,ok := errors.AsType [* daemon.APIError ](err );! ok || failure .StatusCode != 404 { 198t .Fatalf ("missing cancellation: %v" ,err ) 199 } 200} 201 202// Both identified submissions contribute to the same next turn. Targeted 203// cancellation must leave that shared turn running, including anonymous input. 204func TestSubmissionCancellationLeavesSharedRunningTurn (t * testing.T ) { 205firstEntered ,firstRelease := make (chan struct {}),make (chan struct {}) 206sharedEntered ,sharedRelease := make (chan struct {}),make (chan struct {}) 207firstReleased ,sharedReleased := false ,false 208defer func () { 209if ! firstReleased { 210close (firstRelease ) 211 } 212if ! sharedReleased { 213close (sharedRelease ) 214 } 215 }() 216profile := providerRoute (t ,func (request map [string ]any )string { 217if lastUserText (request )== "blocking turn" { 218close (firstEntered ) 219<- firstRelease 220 }else { 221close (sharedEntered ) 222<- sharedRelease 223 } 224return echoReply (request ) 225 }) 226t .Parallel () 227id := newSession (t ,t .TempDir ()) 228client := daemon .NewChatClient (conn (t ),id ) 229if _ ,err := client .Send (t .Context (),"blocking turn" ,nil );err != nil { 230t .Fatal (err ) 231 } 232select { 233case <- firstEntered : 234case <- time .After (10 * time .Second ): 235t .Fatal ("first turn never started" ) 236 } 237shared ,err := client .Send (t .Context (),"identified contribution" ,nil ) 238if err != nil { 239t .Fatal (err ) 240 } 241if _ ,err := client .Send (t .Context (),"anonymous contribution" ,nil );err != nil { 242t .Fatal (err ) 243 } 244queued ,err := client .Send (t .Context (),"cancel this queued contribution" ,nil ) 245if err != nil { 246t .Fatal (err ) 247 } 248outcome ,cancelErr := client .CancelSubmission (t .Context (),queued .OperationID ) 249if cancelErr != nil || outcome != "cancelled_queued" { 250t .Fatalf ("queued cancellation = %q %v" ,outcome ,cancelErr ) 251 } 252outcome ,cancelErr = client .CancelSubmission (t .Context (),queued .OperationID ) 253if cancelErr != nil || outcome != "not_pending" { 254t .Fatalf ("repeated queued cancellation = %q %v" ,outcome ,cancelErr ) 255 } 256close (firstRelease ) 257firstReleased = true 258select { 259case <- sharedEntered : 260case <- time .After (10 * time .Second ): 261t .Fatal ("shared turn never started" ) 262 } 263for range 2 { 264outcome ,err := client .CancelSubmission (t .Context (),shared .OperationID ) 265if err != nil || outcome != "shared_running" { 266t .Fatalf ("shared cancellation = %q, %v" ,outcome ,err ) 267 } 268 } 269status ,err := client .GetStatus (t .Context ()) 270if err != nil || status .Idle { 271t .Fatalf ("shared turn stopped: %+v %v" ,status ,err ) 272 } 273close (sharedRelease ) 274sharedReleased = true 275waitIdle (t ,id ,profile ,2 ) 276requests := suite .provider .requests (profile ) 277encoded ,marshalErr := json .Marshal (requests [1 ]) 278if marshalErr != nil || strings .Contains (string (encoded ),"cancel this queued contribution" )|| ! strings .Contains (string (encoded ),"identified contribution" )|| ! strings .Contains (string (encoded ),"anonymous contribution" ) { 279t .Fatalf ("queued cancellation damaged sibling inputs: %s %v" ,encoded ,marshalErr ) 280 } 281} 282 283// Cancellation races turn admission and completion. Once that target settles, 284// repeating cancellation cannot stop a later exclusively owned turn. 285func TestSubmissionCancellationRacesStartAndCompletion (t * testing.T ) { 286firstEntered ,firstRelease := make (chan struct {}),make (chan struct {}) 287laterEntered ,laterRelease := make (chan struct {}),make (chan struct {}) 288firstReleased ,laterReleased := false ,false 289defer func () { 290if ! firstReleased { 291close (firstRelease ) 292 } 293if ! laterReleased { 294close (laterRelease ) 295 } 296 }() 297profile := providerRoute (t ,func (request map [string ]any )string { 298switch lastUserText (request ) { 299case "hold first turn" : 300close (firstEntered ) 301<- firstRelease 302case "hold later turn" : 303close (laterEntered ) 304<- laterRelease 305 } 306return echoReply (request ) 307 }) 308t .Parallel () 309id := newSession (t ,t .TempDir ()) 310client := daemon .NewChatClient (conn (t ),id ) 311if _ ,err := client .Send (t .Context (),"hold first turn" ,nil );err != nil { 312t .Fatal (err ) 313 } 314select { 315case <- firstEntered : 316case <- time .After (10 * time .Second ): 317t .Fatal ("first turn never started" ) 318 } 319target ,err := client .Send (t .Context (),"short target turn" ,nil ) 320if err != nil { 321t .Fatal (err ) 322 } 323submissionID := target .OperationID 324type cancellation struct { 325err error 326outcome string 327 } 328result := make (chan cancellation ,1 ) 329start := make (chan struct {}) 330go func () { 331<- start 332outcome ,err := client .CancelSubmission (t .Context (),submissionID ) 333result <- cancellation {err :err ,outcome :outcome } 334 }() 335close (start ) 336close (firstRelease ) 337firstReleased = true 338var cancelled cancellation 339select { 340case cancelled = <- result : 341case <- time .After (10 * time .Second ): 342t .Fatal ("racing cancellation did not return" ) 343 } 344if cancelled .err != nil { 345t .Fatal (cancelled .err ) 346 } 347switch cancelled .outcome { 348case "cancelled_queued" ,"interrupt_requested" ,"not_pending" : 349default : 350t .Fatalf ("unexpected racing cancellation outcome %q" ,cancelled .outcome ) 351 } 352waitIdle (t ,id ,profile ,1 ) 353if _ ,err := client .Send (t .Context (),"hold later turn" ,nil );err != nil { 354t .Fatal (err ) 355 } 356select { 357case <- laterEntered : 358case <- time .After (10 * time .Second ): 359t .Fatal ("later turn never started" ) 360 } 361outcome ,err := client .CancelSubmission (t .Context (),submissionID ) 362if err != nil || outcome != "not_pending" { 363t .Fatalf ("settled target cancellation = %q %v" ,outcome ,err ) 364 } 365status ,err := client .GetStatus (t .Context ()) 366if err != nil || status .Idle { 367t .Fatalf ("old cancellation stopped a later turn: %+v %v" ,status ,err ) 368 } 369close (laterRelease ) 370laterReleased = true 371waitIdle (t ,id ,profile ,2 ) 372}