3e037ba fix(session): harden resume lookup cancellation · zut

M packages/agent/cli.go +11 -7


@@ -899,7 +899,7 @@ func runPrintMode(ctx context.Context, args Args, version string) error {
899	899		}
900	900		}()
901	901		wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
902		-	sess, err := openOrCreateSession(args, r, ag, version)
	902	+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
903	903		if err != nil {
904	904		return err
905	905		}

@@ -972,7 +972,7 @@ func runStreamMode(ctx context.Context, args Args, version string) error {
972	972		}
973	973		}()
974	974		wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
975		-	sess, err := openOrCreateSession(args, r, ag, version)
	975	+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
976	976		if err != nil {
977	977		return err
978	978		}

@@ -1084,7 +1084,7 @@ func runJSONMode(ctx context.Context, args Args, version string) error {
1084	1084		}
1085	1085		}()
1086	1086		wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
1087		-	sess, err := openOrCreateSession(args, r, ag, version)
	1087	+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
1088	1088		if err != nil {
1089	1089		return err
1090	1090		}

@@ -1716,7 +1716,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1716	1716		var sessionTitlePending bool
1717	1717		if !args.NoSess && ag != nil {
1718	1718		var sessErr error
1719		-	sess, sessErr = openOrCreateSession(args, r, ag, version)
	1719	+	sess, sessErr = openOrCreateSession(ctx, args, r, ag, version)
1720	1720		if sessErr != nil {
1721	1721		return sessErr
1722	1722		}

@@ -2599,10 +2599,14 @@ func resumableSessionID(root string, sess *core.Session) string {
2599	2599		if sess == nil \|\| sess.ID == "" \|\| sess.Path == "" {
2600	2600		return ""
2601	2601		}
2602		-	path, err := core.FindManagedSessionByID(root, sess.ID)
	2602	+	path, err := core.FindManagedSessionByID(context.Background(), root, sess.ID)
2603	2603		if err != nil \|\| !sameSessionPath(path, sess.Path) {
2604	2604		return ""
2605	2605		}
	2606	+	snapshot, err := core.ReadSessionSnapshot(path)
	2607	+	if err != nil \|\| len(snapshot.Messages) == 0 {
	2608	+	return ""
	2609	+	}
2606	2610		return sess.ID
2607	2611		}
2608	2612

@@ -2628,7 +2632,7 @@ func agentSessionsRoot(root string, args Args) string {
2628	2632
2629	2633		// openOrCreateSession returns a session for the run. sess may be nil
2630	2634		// with a nil error if session persistence is disabled.
2631		-	func openOrCreateSession(args Args, r Resolved, ag core.Agent, version string) (core.Session, error) {
	2635	+	func openOrCreateSession(ctx context.Context, args Args, r Resolved, ag core.Agent, version string) (core.Session, error) {
2632	2636		if args.NoSess {
2633	2637		return nil, nil
2634	2638		}

@@ -2667,7 +2671,7 @@ func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string)
2667	2671		}
2668	2672		case args.Resume:
2669	2673		if args.ResumeSessionID != "" {
2670		-	picked, lookupErr := core.FindManagedSessionByID(ZutHome(), args.ResumeSessionID)
	2674	+	picked, lookupErr := core.FindManagedSessionByID(ctx, ZutHome(), args.ResumeSessionID)
2671	2675		if lookupErr != nil {
2672	2676		return nil, lookupErr
2673	2677		}

 		}
 	}()
 	wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
-	sess, err := openOrCreateSession(args, r, ag, version)
+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
 	if err != nil {
 		return err
 	}
 		}
 	}()
 	wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
-	sess, err := openOrCreateSession(args, r, ag, version)
+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
 	if err != nil {
 		return err
 	}
 		}
 	}()
 	wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
-	sess, err := openOrCreateSession(args, r, ag, version)
+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
 	if err != nil {
 		return err
 	}
 	var sessionTitlePending bool
 	if !args.NoSess && ag != nil {
 		var sessErr error
-		sess, sessErr = openOrCreateSession(args, r, ag, version)
+		sess, sessErr = openOrCreateSession(ctx, args, r, ag, version)
 		if sessErr != nil {
 			return sessErr
 		}
 	if sess == nil || sess.ID == "" || sess.Path == "" {
 		return ""
 	}
-	path, err := core.FindManagedSessionByID(root, sess.ID)
+	path, err := core.FindManagedSessionByID(context.Background(), root, sess.ID)
 	if err != nil || !sameSessionPath(path, sess.Path) {
 		return ""
 	}
+	snapshot, err := core.ReadSessionSnapshot(path)
+	if err != nil || len(snapshot.Messages) == 0 {
+		return ""
+	}
 	return sess.ID
 }
 // openOrCreateSession returns a session for the run. sess may be nil
 // with a nil error if session persistence is disabled.
-func openOrCreateSession(args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
+func openOrCreateSession(ctx context.Context, args Args, r Resolved, ag *core.Agent, version string) (*core.Session, error) {
 	if args.NoSess {
 		return nil, nil
 	}
 		}
 	case args.Resume:
 		if args.ResumeSessionID != "" {
-			picked, lookupErr := core.FindManagedSessionByID(ZutHome(), args.ResumeSessionID)
+			picked, lookupErr := core.FindManagedSessionByID(ctx, ZutHome(), args.ResumeSessionID)
 			if lookupErr != nil {
 				return nil, lookupErr
 			}

M packages/agent/cli_session_test.go +57 -4


@@ -40,7 +40,7 @@ func TestOpenOrCreateSessionResumesByUUIDAcrossCWDAndAgentStores(t *testing.T) {
40	40		}
41	41
42	42		ag := core.NewAgent(nil, "model", "", nil)
43		-	got, err := openOrCreateSession(Args{
	43	+	got, err := openOrCreateSession(context.Background(), Args{
44	44		CWD: otherCWD,
45	45		Resume: true,
46	46		ResumeSessionID: id,

@@ -75,7 +75,7 @@ func TestOpenOrCreateSessionResumesByUUIDAcrossCWDAndAgentStores(t *testing.T) {
75	75		}
76	76
77	77		namedAgent := core.NewAgent(nil, "model", "", nil)
78		-	namedGot, err := openOrCreateSession(Args{
	78	+	namedGot, err := openOrCreateSession(context.Background(), Args{
79	79		CWD: otherCWD,
80	80		Resume: true,
81	81		ResumeSessionID: namedID,

@@ -131,13 +131,66 @@ func TestResumeHintSuppressesEmptyAndArbitrarySessions(t *testing.T) {
131	131		}
132	132		}
133	133
	134	+	func TestResumeHintSuppressesPersistedEmptyExplicitUUIDSession(t *testing.T) {
	135	+	root := t.TempDir()
	136	+	cwd := t.TempDir()
	137	+	otherCWD := t.TempDir()
	138	+	t.Setenv("ZUT_HOME", root)
	139	+
	140	+	dir := core.SessionsDir(root, cwd)
	141	+	if err := os.MkdirAll(dir, 0o755); err != nil {
	142	+	t.Fatal(err)
	143	+	}
	144	+	id := uuid.NewString()
	145	+	metaLine := struct {
	146	+	Type string `json:"type"`
	147	+	Meta core.SessionMeta `json:"meta"`
	148	+	}{
	149	+	Type: "meta",
	150	+	Meta: core.SessionMeta{
	151	+	ID: id,
	152	+	CWD: cwd,
	153	+	Provider: "provider",
	154	+	Model: "model",
	155	+	Started: time.Now().UTC(),
	156	+	Version: "test",
	157	+	},
	158	+	}
	159	+	data, err := json.Marshal(metaLine)
	160	+	if err != nil {
	161	+	t.Fatal(err)
	162	+	}
	163	+	path := filepath.Join(dir, "persisted-empty.jsonl")
	164	+	if err := os.WriteFile(path, append(data, '\n'), 0o644); err != nil {
	165	+	t.Fatal(err)
	166	+	}
	167	+
	168	+	sess, err := openOrCreateSession(context.Background(), Args{
	169	+	CWD: otherCWD,
	170	+	Resume: true,
	171	+	ResumeSessionID: id,
	172	+	}, Resolved{Provider: "provider", Model: "model"}, core.NewAgent(nil, "model", "", nil), "test")
	173	+	if err != nil {
	174	+	t.Fatal(err)
	175	+	}
	176	+	if sess == nil \|\| sess.Path != path {
	177	+	t.Fatalf("resumed persisted empty session = %#v, want path %q", sess, path)
	178	+	}
	179	+	if got := resumableSessionID(root, sess); got != "" {
	180	+	t.Fatalf("persisted empty session hint ID = %q, want empty", got)
	181	+	}
	182	+	if err := sess.Close(); err != nil {
	183	+	t.Fatal(err)
	184	+	}
	185	+	}
	186	+
134	187		func TestOpenOrCreateSessionUnknownUUIDDoesNotCreateSession(t *testing.T) {
135	188		root := t.TempDir()
136	189		cwd := t.TempDir()
137	190		t.Setenv("ZUT_HOME", root)
138	191		id := uuid.NewString()
139	192
140		-	_, err := openOrCreateSession(Args{
	193	+	_, err := openOrCreateSession(context.Background(), Args{
141	194		CWD: cwd,
142	195		Resume: true,
143	196		ResumeSessionID: id,

@@ -164,7 +217,7 @@ func TestOpenOrCreateSessionUUIDLookupReturnsMetadataErrors(t *testing.T) {
164	217		t.Fatal(err)
165	218		}
166	219
167		-	_, err := openOrCreateSession(Args{
	220	+	_, err := openOrCreateSession(context.Background(), Args{
168	221		CWD: cwd,
169	222		Resume: true,
170	223		ResumeSessionID: uuid.NewString(),


@@ -40,7 +40,7 @@ func TestOpenOrCreateSessionResumesByUUIDAcrossCWDAndAgentStores(t *testing.T) {
40		}	40		}
41			41
42		ag := core.NewAgent(nil, "model", "", nil)	42		ag := core.NewAgent(nil, "model", "", nil)
43	-	got, err := openOrCreateSession(Args{	43	+	got, err := openOrCreateSession(context.Background(), Args{
44		CWD: otherCWD,	44		CWD: otherCWD,
45		Resume: true,	45		Resume: true,
46		ResumeSessionID: id,	46		ResumeSessionID: id,

@@ -75,7 +75,7 @@ func TestOpenOrCreateSessionResumesByUUIDAcrossCWDAndAgentStores(t *testing.T) {
75		}	75		}
76			76
77		namedAgent := core.NewAgent(nil, "model", "", nil)	77		namedAgent := core.NewAgent(nil, "model", "", nil)
78	-	namedGot, err := openOrCreateSession(Args{	78	+	namedGot, err := openOrCreateSession(context.Background(), Args{
79		CWD: otherCWD,	79		CWD: otherCWD,
80		Resume: true,	80		Resume: true,
81		ResumeSessionID: namedID,	81		ResumeSessionID: namedID,

@@ -131,13 +131,66 @@ func TestResumeHintSuppressesEmptyAndArbitrarySessions(t *testing.T) {
131		}	131		}
132		}	132		}
133			133
			134	+	func TestResumeHintSuppressesPersistedEmptyExplicitUUIDSession(t *testing.T) {
			135	+	root := t.TempDir()
			136	+	cwd := t.TempDir()
			137	+	otherCWD := t.TempDir()
			138	+	t.Setenv("ZUT_HOME", root)
			139	+
			140	+	dir := core.SessionsDir(root, cwd)
			141	+	if err := os.MkdirAll(dir, 0o755); err != nil {
			142	+	t.Fatal(err)
			143	+	}
			144	+	id := uuid.NewString()
			145	+	metaLine := struct {
			146	+	Type string `json:"type"`
			147	+	Meta core.SessionMeta `json:"meta"`
			148	+	}{
			149	+	Type: "meta",
			150	+	Meta: core.SessionMeta{
			151	+	ID: id,
			152	+	CWD: cwd,
			153	+	Provider: "provider",
			154	+	Model: "model",
			155	+	Started: time.Now().UTC(),
			156	+	Version: "test",
			157	+	},
			158	+	}
			159	+	data, err := json.Marshal(metaLine)
			160	+	if err != nil {
			161	+	t.Fatal(err)
			162	+	}
			163	+	path := filepath.Join(dir, "persisted-empty.jsonl")
			164	+	if err := os.WriteFile(path, append(data, '\n'), 0o644); err != nil {
			165	+	t.Fatal(err)
			166	+	}
			167	+
			168	+	sess, err := openOrCreateSession(context.Background(), Args{
			169	+	CWD: otherCWD,
			170	+	Resume: true,
			171	+	ResumeSessionID: id,
			172	+	}, Resolved{Provider: "provider", Model: "model"}, core.NewAgent(nil, "model", "", nil), "test")
			173	+	if err != nil {
			174	+	t.Fatal(err)
			175	+	}
			176	+	if sess == nil \|\| sess.Path != path {
			177	+	t.Fatalf("resumed persisted empty session = %#v, want path %q", sess, path)
			178	+	}
			179	+	if got := resumableSessionID(root, sess); got != "" {
			180	+	t.Fatalf("persisted empty session hint ID = %q, want empty", got)
			181	+	}
			182	+	if err := sess.Close(); err != nil {
			183	+	t.Fatal(err)
			184	+	}
			185	+	}
			186	+
134		func TestOpenOrCreateSessionUnknownUUIDDoesNotCreateSession(t *testing.T) {	187		func TestOpenOrCreateSessionUnknownUUIDDoesNotCreateSession(t *testing.T) {
135		root := t.TempDir()	188		root := t.TempDir()
136		cwd := t.TempDir()	189		cwd := t.TempDir()
137		t.Setenv("ZUT_HOME", root)	190		t.Setenv("ZUT_HOME", root)
138		id := uuid.NewString()	191		id := uuid.NewString()
139			192
140	-	_, err := openOrCreateSession(Args{	193	+	_, err := openOrCreateSession(context.Background(), Args{
141		CWD: cwd,	194		CWD: cwd,
142		Resume: true,	195		Resume: true,
143		ResumeSessionID: id,	196		ResumeSessionID: id,

@@ -164,7 +217,7 @@ func TestOpenOrCreateSessionUUIDLookupReturnsMetadataErrors(t *testing.T) {
164		t.Fatal(err)	217		t.Fatal(err)
165		}	218		}
166			219
167	-	_, err := openOrCreateSession(Args{	220	+	_, err := openOrCreateSession(context.Background(), Args{
168		CWD: cwd,	221		CWD: cwd,
169		Resume: true,	222		Resume: true,
170		ResumeSessionID: uuid.NewString(),	223		ResumeSessionID: uuid.NewString(),

M packages/agent/subagent_worker.go +1 -1


@@ -70,7 +70,7 @@ func runSubagentWorkerMode(ctx context.Context, args Args, version string) error
70	70		}
71	71		}()
72	72		wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
73		-	sess, err := openOrCreateSession(args, r, ag, version)
	73	+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
74	74		if err != nil {
75	75		return err
76	76		}


@@ -70,7 +70,7 @@ func runSubagentWorkerMode(ctx context.Context, args Args, version string) error
70		}	70		}
71		}()	71		}()
72		wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)	72		wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
73	-	sess, err := openOrCreateSession(args, r, ag, version)	73	+	sess, err := openOrCreateSession(ctx, args, r, ag, version)
74		if err != nil {	74		if err != nil {
75		return err	75		return err
76		}	76		}

M packages/core/session_portable.go +28 -5


@@ -3,6 +3,7 @@ package core
3	3		import (
4	4		"bufio"
5	5		"bytes"
	6	+	"context"
6	7		"encoding/json"
7	8		"errors"
8	9		"fmt"

@@ -803,6 +804,13 @@ func linkSessionTreeNodes(nodes map[string]TreeNode, order []string) []TreeNod
803	804		// small projection; full snapshot validation remains the responsibility of
804	805		// resume, export, branching, and tree preflight callers.
805	806		func scanSessionMeta(path string) (meta SessionMeta, err error) {
	807	+	return scanSessionMetaContext(context.Background(), path)
	808	+	}
	809	+
	810	+	func scanSessionMetaContext(ctx context.Context, path string) (meta SessionMeta, err error) {
	811	+	if err := contextErr(ctx); err != nil {
	812	+	return SessionMeta{}, err
	813	+	}
806	814		f, err := os.Open(path)
807	815		if err != nil {
808	816		return SessionMeta{}, err

@@ -810,7 +818,7 @@ func scanSessionMeta(path string) (meta SessionMeta, err error) {
810	818		defer f.Close()
811	819
812	820		var sawMeta bool
813		-	err = forEachStrictJSONLLine(f, func(line []byte, lineNo int) error {
	821	+	err = forEachStrictJSONLLineContext(ctx, f, func(line []byte, lineNo int) error {
814	822		var head sessionLineHead
815	823		if err := json.Unmarshal(line, &head); err != nil {
816	824		return fmt.Errorf("line %d: invalid JSON: %w", lineNo, err)

@@ -886,7 +894,7 @@ func FindSessionByID(root, cwd, id string) string {
886	894		// session stores, independent of cwd. A nil error and empty path mean no
887	895		// matching session was found. Read, storage, metadata, and duplicate-ID
888	896		// failures are returned to the caller instead of being reported as a miss.
889		-	func FindManagedSessionByID(root, id string) (string, error) {
	897	+	func FindManagedSessionByID(ctx context.Context, root, id string) (string, error) {
890	898		if strings.TrimSpace(root) == "" {
891	899		return "", errors.New("find managed session: root is empty")
892	900		}

@@ -910,10 +918,16 @@ func FindManagedSessionByID(root, id string) (string, error) {
910	918
911	919		var matches []string
912	920		for _, store := range stores {
913		-	if err := collectManagedSessionMatches(store, id, &matches); err != nil {
	921	+	if err := contextErr(ctx); err != nil {
	922	+	return "", err
	923	+	}
	924	+	if err := collectManagedSessionMatches(ctx, store, id, &matches); err != nil {
914	925		return "", err
915	926		}
916	927		}
	928	+	if err := contextErr(ctx); err != nil {
	929	+	return "", err
	930	+	}
917	931		if len(matches) == 0 {
918	932		return "", nil
919	933		}

@@ -923,7 +937,10 @@ func FindManagedSessionByID(root, id string) (string, error) {
923	937		return matches[0], nil
924	938		}
925	939
926		-	func collectManagedSessionMatches(root, id string, matches *[]string) error {
	940	+	func collectManagedSessionMatches(ctx context.Context, root, id string, matches *[]string) error {
	941	+	if err := contextErr(ctx); err != nil {
	942	+	return err
	943	+	}
927	944		dir := filepath.Join(root, "sessions")
928	945		entries, err := os.ReadDir(dir)
929	946		if errors.Is(err, os.ErrNotExist) {

@@ -933,6 +950,9 @@ func collectManagedSessionMatches(root, id string, matches *[]string) error {
933	950		return fmt.Errorf("find managed session: read %q: %w", dir, err)
934	951		}
935	952		for _, bucket := range entries {
	953	+	if err := contextErr(ctx); err != nil {
	954	+	return err
	955	+	}
936	956		if !bucket.IsDir() \|\| bucket.Name() == "agents" \|\| !isSessionBucketName(bucket.Name()) {
937	957		continue
938	958		}

@@ -942,11 +962,14 @@ func collectManagedSessionMatches(root, id string, matches *[]string) error {
942	962		return fmt.Errorf("find managed session: read %q: %w", bucketPath, err)
943	963		}
944	964		for _, file := range files {
	965	+	if err := contextErr(ctx); err != nil {
	966	+	return err
	967	+	}
945	968		if file.IsDir() \|\| !strings.HasSuffix(file.Name(), ".jsonl") {
946	969		continue
947	970		}
948	971		path := filepath.Join(bucketPath, file.Name())
949		-	meta, err := scanSessionMeta(path)
	972	+	meta, err := scanSessionMetaContext(ctx, path)
950	973		if err != nil {
951	974		return fmt.Errorf("find managed session: read metadata %q: %w", path, err)
952	975		}

 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 // small projection; full snapshot validation remains the responsibility of
 // resume, export, branching, and tree preflight callers.
 func scanSessionMeta(path string) (meta SessionMeta, err error) {
+	return scanSessionMetaContext(context.Background(), path)
+}
+func scanSessionMetaContext(ctx context.Context, path string) (meta SessionMeta, err error) {
+	if err := contextErr(ctx); err != nil {
+		return SessionMeta{}, err
+	}
 	f, err := os.Open(path)
 	if err != nil {
 		return SessionMeta{}, err
 	defer f.Close()
 	var sawMeta bool
-	err = forEachStrictJSONLLine(f, func(line []byte, lineNo int) error {
+	err = forEachStrictJSONLLineContext(ctx, f, func(line []byte, lineNo int) error {
 		var head sessionLineHead
 		if err := json.Unmarshal(line, &head); err != nil {
 			return fmt.Errorf("line %d: invalid JSON: %w", lineNo, err)
 // session stores, independent of cwd. A nil error and empty path mean no
 // matching session was found. Read, storage, metadata, and duplicate-ID
 // failures are returned to the caller instead of being reported as a miss.
-func FindManagedSessionByID(root, id string) (string, error) {
+func FindManagedSessionByID(ctx context.Context, root, id string) (string, error) {
 	if strings.TrimSpace(root) == "" {
 		return "", errors.New("find managed session: root is empty")
 	}
 	var matches []string
 	for _, store := range stores {
-		if err := collectManagedSessionMatches(store, id, &matches); err != nil {
+		if err := contextErr(ctx); err != nil {
+			return "", err
+		}
+		if err := collectManagedSessionMatches(ctx, store, id, &matches); err != nil {
 			return "", err
 		}
 	}
+	if err := contextErr(ctx); err != nil {
+		return "", err
+	}
 	if len(matches) == 0 {
 		return "", nil
 	}
 	return matches[0], nil
 }
-func collectManagedSessionMatches(root, id string, matches *[]string) error {
+func collectManagedSessionMatches(ctx context.Context, root, id string, matches *[]string) error {
+	if err := contextErr(ctx); err != nil {
+		return err
+	}
 	dir := filepath.Join(root, "sessions")
 	entries, err := os.ReadDir(dir)
 	if errors.Is(err, os.ErrNotExist) {
 		return fmt.Errorf("find managed session: read %q: %w", dir, err)
 	}
 	for _, bucket := range entries {
+		if err := contextErr(ctx); err != nil {
+			return err
+		}
 		if !bucket.IsDir() || bucket.Name() == "agents" || !isSessionBucketName(bucket.Name()) {
 			continue
 		}
 			return fmt.Errorf("find managed session: read %q: %w", bucketPath, err)
 		}
 		for _, file := range files {
+			if err := contextErr(ctx); err != nil {
+				return err
+			}
 			if file.IsDir() || !strings.HasSuffix(file.Name(), ".jsonl") {
 				continue
 			}
 			path := filepath.Join(bucketPath, file.Name())
-			meta, err := scanSessionMeta(path)
+			meta, err := scanSessionMetaContext(ctx, path)
 			if err != nil {
 				return fmt.Errorf("find managed session: read metadata %q: %w", path, err)
 			}

M packages/core/session_portable_test.go +13 -2

@@ -1,7 +1,9 @@
1	1		package core
2	2
3	3		import (
	4	+	"context"
4	5		"encoding/json"
	6	+	"errors"
5	7		"os"
6	8		"path/filepath"
7	9		"strings"

@@ -27,7 +29,7 @@ func TestFindManagedSessionByIDStrictScopesAllManagedStoresAndReturnsReadErrors(
27	29		t.Fatal(err)
28	30		}
29	31
30		-	got, err := FindManagedSessionByID(root, session.ID)
	32	+	got, err := FindManagedSessionByID(context.Background(), root, session.ID)
31	33		if err != nil {
32	34		t.Fatalf("FindManagedSessionByID: %v", err)
33	35		}

@@ -39,13 +41,22 @@ func TestFindManagedSessionByIDStrictScopesAllManagedStoresAndReturnsReadErrors(
39	41		if err := os.WriteFile(badPath, []byte("not json\n"), 0o644); err != nil {
40	42		t.Fatal(err)
41	43		}
42		-	if _, err := FindManagedSessionByID(root, "missing"); err == nil {
	44	+	if _, err := FindManagedSessionByID(context.Background(), root, "missing"); err == nil {
43	45		t.Fatal("malformed managed metadata was reported as not found")
44	46		} else if !strings.Contains(err.Error(), "metadata") {
45	47		t.Fatalf("malformed metadata error = %v, want metadata context", err)
46	48		}
47	49		}
48	50
	51	+	func TestFindManagedSessionByIDHonorsCancellation(t *testing.T) {
	52	+	ctx, cancel := context.WithCancel(context.Background())
	53	+	cancel()
	54	+
	55	+	if _, err := FindManagedSessionByID(ctx, t.TempDir(), "session-id"); !errors.Is(err, context.Canceled) {
	56	+	t.Fatalf("FindManagedSessionByID error = %v, want context.Canceled", err)
	57	+	}
	58	+	}
	59	+
49	60		// TestSessionExportImportRoundTrip writes a few messages to a live
50	61		// session, exports it, imports the export under a different cwd,
51	62		// and verifies OpenSession on the imported file yields the same

 package core
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"os"
 	"path/filepath"
 	"strings"
 		t.Fatal(err)
 	}
-	got, err := FindManagedSessionByID(root, session.ID)
+	got, err := FindManagedSessionByID(context.Background(), root, session.ID)
 	if err != nil {
 		t.Fatalf("FindManagedSessionByID: %v", err)
 	}
 	if err := os.WriteFile(badPath, []byte("not json\n"), 0o644); err != nil {
 		t.Fatal(err)
 	}
-	if _, err := FindManagedSessionByID(root, "missing"); err == nil {
+	if _, err := FindManagedSessionByID(context.Background(), root, "missing"); err == nil {
 		t.Fatal("malformed managed metadata was reported as not found")
 	} else if !strings.Contains(err.Error(), "metadata") {
 		t.Fatalf("malformed metadata error = %v, want metadata context", err)
 	}
 }
+func TestFindManagedSessionByIDHonorsCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if _, err := FindManagedSessionByID(ctx, t.TempDir(), "session-id"); !errors.Is(err, context.Canceled) {
+		t.Fatalf("FindManagedSessionByID error = %v, want context.Canceled", err)
+	}
+}
 // TestSessionExportImportRoundTrip writes a few messages to a live
 // session, exports it, imports the export under a different cwd,
 // and verifies OpenSession on the imported file yields the same