@@ -2,6 +2,7 @@ package tui
22
33import (
44 "encoding/json"
5+ "errors"
56 "net/http"
67 "net/http/httptest"
78 "strings"
@@ -16,24 +17,62 @@ import (
1617
1718// wired builds a Model backed by a live in-process odek-serve stand-in.
1819func wired (t * testing.T ) * Model {
20+ t .Helper ()
21+ return standIn (t , "tok" )
22+ }
23+
24+ // standIn builds a Model against an in-process odek-serve stand-in that
25+ // optionally enforces the per-instance WS token (cookie or X-Odek-Ws-Token
26+ // header) on /ws and /api/*, mirroring odek serve; an empty token disables
27+ // enforcement like odek versions that predate enforced auth.
28+ func standIn (t * testing.T , token string ) * Model {
1929 t .Helper ()
2030 t .Setenv ("HOME" , t .TempDir ())
2131
22- mux := http .NewServeMux ()
23- mux .Handle ("/ws" , ws .Handler (func (c * ws.Conn ) {
24- for {
25- var d []byte
26- if err := ws .Message .Receive (c , & d ); err != nil {
32+ // serveTokenOK mirrors odek serve's validateServeToken: the token is
33+ // accepted via the odek_ws_token cookie or the X-Odek-Ws-Token header.
34+ serveTokenOK := func (r * http.Request ) bool {
35+ if token == "" {
36+ return true
37+ }
38+ if c , err := r .Cookie ("odek_ws_token" ); err == nil && c .Value == token {
39+ return true
40+ }
41+ return r .Header .Get ("X-Odek-Ws-Token" ) == token
42+ }
43+ guard := func (h http.HandlerFunc ) http.HandlerFunc {
44+ return func (w http.ResponseWriter , r * http.Request ) {
45+ if ! serveTokenOK (r ) {
46+ w .WriteHeader (http .StatusForbidden )
2747 return
2848 }
29- _ = ws .JSON .Send (c , map [string ]any {"type" : "session" , "session_id" : "s1" , "auth_token" : "a1" , "model" : "m" })
30- _ = ws .Message .Send (c , `{"type":"done","latency":1}` )
49+ h (w , r )
3150 }
32- }))
33- mux .HandleFunc ("/api/sessions" , func (w http.ResponseWriter , r * http.Request ) {
34- json .NewEncoder (w ).Encode ([]client.Session {{ID : "s1" , Task : "first task" , Turns : 1 , UpdatedAt : time .Now ()}})
51+ }
52+
53+ mux := http .NewServeMux ()
54+ mux .Handle ("/ws" , ws.Server {
55+ Handshake : func (cfg * ws.Config , r * http.Request ) error {
56+ if ! serveTokenOK (r ) {
57+ return errors .New ("missing or invalid server token" )
58+ }
59+ return nil
60+ },
61+ Handler : ws .Handler (func (c * ws.Conn ) {
62+ for {
63+ var d []byte
64+ if err := ws .Message .Receive (c , & d ); err != nil {
65+ return
66+ }
67+ _ = ws .JSON .Send (c , map [string ]any {"type" : "session" , "session_id" : "s1" , "auth_token" : "a1" , "model" : "m" })
68+ _ = ws .Message .Send (c , `{"type":"done","latency":1}` )
69+ }
70+ }),
3571 })
36- mux .HandleFunc ("/api/sessions/" , func (w http.ResponseWriter , r * http.Request ) {
72+ mux .HandleFunc ("/api/sessions" , guard (func (w http.ResponseWriter , r * http.Request ) {
73+ json .NewEncoder (w ).Encode ([]client.Session {{ID : "s1" , Task : "first task" , Turns : 1 , UpdatedAt : time .Now ()}})
74+ }))
75+ mux .HandleFunc ("/api/sessions/" , guard (func (w http.ResponseWriter , r * http.Request ) {
3776 if r .Method == http .MethodDelete {
3877 w .WriteHeader (http .StatusNoContent )
3978 return
@@ -47,20 +86,20 @@ func wired(t *testing.T) *Model {
4786 {Role : "assistant" , Content : "" }, // skipped (empty)
4887 },
4988 })
50- })
51- mux .HandleFunc ("/api/models" , func (w http.ResponseWriter , r * http.Request ) {
89+ }))
90+ mux .HandleFunc ("/api/models" , guard ( func (w http.ResponseWriter , r * http.Request ) {
5291 json .NewEncoder (w ).Encode ([]client.ModelInfo {{ID : "m1" , Description : "one" , Current : true }})
53- })
54- mux .HandleFunc ("/api/resources" , func (w http.ResponseWriter , r * http.Request ) {
92+ }))
93+ mux .HandleFunc ("/api/resources" , guard ( func (w http.ResponseWriter , r * http.Request ) {
5594 json .NewEncoder (w ).Encode ([]client.Resource {{ID : "@main.go" , Type : "file" , Label : "main.go" , Detail : "1 KB" }})
56- })
57- mux .HandleFunc ("/api/cancel" , func (w http.ResponseWriter , r * http.Request ) {
95+ }))
96+ mux .HandleFunc ("/api/cancel" , guard ( func (w http.ResponseWriter , r * http.Request ) {
5897 w .WriteHeader (http .StatusNoContent )
59- })
98+ }))
6099
61100 srv := httptest .NewServer (mux )
62101 t .Cleanup (srv .Close )
63- cl , err := client .Dial ("ws" + strings .TrimPrefix (srv .URL , "http" )+ "/ws" , srv .URL , srv .URL , "tok" )
102+ cl , err := client .Dial ("ws" + strings .TrimPrefix (srv .URL , "http" )+ "/ws" , srv .URL , srv .URL , token )
64103 if err != nil {
65104 t .Fatalf ("Dial: %v" , err )
66105 }
@@ -71,6 +110,16 @@ func wired(t *testing.T) *Model {
71110 return m
72111}
73112
113+ func TestStandInWithoutTokenEnforcement (t * testing.T ) {
114+ // Empty token = an odek that predates enforced auth: the stand-in serves
115+ // /ws and /api/* without any token checks.
116+ m := standIn (t , "" )
117+ m .Update (exec (m .openModels ()))
118+ if len (m .models ) != 1 {
119+ t .Fatalf ("models against an unenforced stand-in: %d" , len (m .models ))
120+ }
121+ }
122+
74123func exec (cmd tea.Cmd ) tea.Msg {
75124 if cmd == nil {
76125 return nil
@@ -357,4 +406,8 @@ func TestElapsed(t *testing.T) {
357406 if m .elapsed () == "" {
358407 t .Error ("elapsed should be non-empty during a run" )
359408 }
409+ // The live badge ticks in whole seconds (tenths would flicker at 12fps).
410+ if strings .Contains (m .elapsed (), "." ) {
411+ t .Errorf ("elapsed should not show tenths: %q" , m .elapsed ())
412+ }
360413}
0 commit comments