diff --git a/docs/design/agentcube-proposal.md b/docs/design/agentcube-proposal.md index 1bfa2c022..ed7cc3bbf 100644 --- a/docs/design/agentcube-proposal.md +++ b/docs/design/agentcube-proposal.md @@ -277,6 +277,7 @@ POST /v1/agent-runtime { "namespace": "string", // agentruntime CR namespace, required "name": "string", // agentruntime CR name, required + "ttl": 3600 // requested maximum lifetime in seconds, optional } ``` @@ -311,8 +312,15 @@ POST /v1/code-interpreter { "namespace": "string", // codeinterpreter CR namespace, required "name": "string", // codeinterpreter CR name, required + "ttl": 3600 // requested maximum lifetime in seconds, optional } ``` + +For both creation APIs, the effective maximum lifetime is the smaller of the +requested `ttl` and the workload CR's `spec.maxSessionDuration`. If `ttl` is +omitted, `spec.maxSessionDuration` (or its default) is used. The independent +`spec.sessionTimeout` idle limit may reclaim the session earlier. + Response body: ```json diff --git a/docs/design/router-proposal.md b/docs/design/router-proposal.md index 160bed16d..66b7247eb 100644 --- a/docs/design/router-proposal.md +++ b/docs/design/router-proposal.md @@ -135,7 +135,7 @@ The Router uses the Gin framework to provide HTTP services with the following en - Preserves original response status and body **Error Handling:** -- Invalid session ID → `400 Bad Request` +- Missing or expired session ID → `404 Not Found` with code `SESSION_NOT_FOUND` - No entry points → `404 Not Found` - Invalid endpoint → `500 Internal Server Error` - Connection refused → `502 Bad Gateway` (SANDBOX_UNREACHABLE) @@ -170,7 +170,7 @@ The Router uses the Gin framework to provide HTTP services with the following en | Status Code | Scenario | Response Body Example | |-------------|----------|----------------------| -| 400 Bad Request | Invalid session ID | `{"error": "Invalid session id ", "code": "BadRequest"}` | +| 404 Not Found | Missing or expired session ID | `{"error": "session was not found", "code": "SESSION_NOT_FOUND"}` | | 404 Not Found | No entry points found for sandbox | `{"error": "no entry points found for sandbox", "code": "Service not found"}` | | 429 Too Many Requests | Server overloaded (concurrent request limit exceeded) | `{"error": "server overloaded, please try again later", "code": "SERVER_OVERLOADED"}` | diff --git a/docs/devguide/code-interpreter-python-sdk.md b/docs/devguide/code-interpreter-python-sdk.md index a19b600ed..6d75c94d5 100644 --- a/docs/devguide/code-interpreter-python-sdk.md +++ b/docs/devguide/code-interpreter-python-sdk.md @@ -26,7 +26,7 @@ The `CodeInterpreterClient` is the main entry point. You can initialize it direc |-----------|------|---------|-------------| | `name` | `str` | `"simple-codeinterpreter"` | CodeInterpreter CRD template name | | `namespace` | `str` | `"default"` | Kubernetes namespace | -| `ttl` | `int` | `3600` | Session time-to-live (seconds) | +| `ttl` | `Optional[int]` | `None` | Optional requested maximum session lifetime in seconds, capped by `spec.maxSessionDuration` | | `workload_manager_url` | `str` | `None` | Control Plane URL (falls back to env `WORKLOAD_MANAGER_URL`) | | `router_url` | `str` | `None` | Data Plane Router URL (falls back to env `ROUTER_URL`) | | `auth_token` | `str` | `None` | Auth token (falls back to K8s SA token) | @@ -202,3 +202,24 @@ client2 = CodeInterpreterClient(session_id=session_id) client2.run_code("python", "print(open('value.txt').read())") client2.stop() # Cleanup when done ``` + +If a reused session has already been reclaimed, the SDK raises +`SessionNotFoundError` and clears the stale ID. It does not silently replace +the session because files and other sandbox-local state are no longer +available. + +```python +from agentcube import CodeInterpreterClient, SessionNotFoundError + +client = CodeInterpreterClient(session_id=saved_session_id) +try: + client.list_files() +except SessionNotFoundError: + # Decide at the application layer whether starting without old state is safe. + client = CodeInterpreterClient() +``` + +Session reclamation occurs when either the idle `spec.sessionTimeout` or the +effective maximum lifetime is reached. If `ttl` is omitted, the CRD's +`spec.maxSessionDuration` determines the maximum lifetime. Otherwise, the +effective maximum is the smaller of `ttl` and `spec.maxSessionDuration`. diff --git a/pkg/api/errors.go b/pkg/api/errors.go index 6bc6f6547..e60ada1f5 100644 --- a/pkg/api/errors.go +++ b/pkg/api/errors.go @@ -31,6 +31,9 @@ const ( sessionResourceName = "sessions" agentRuntimeResourceName = "agentruntimes" codeInterpreterResourceName = "codeinterpreters" + // SessionNotFoundCode is returned when an invocation references a session + // that is no longer present in the session store. + SessionNotFoundCode = "SESSION_NOT_FOUND" ) var ( @@ -57,6 +60,20 @@ func NewSessionNotFoundError(sessionID string) error { return apierrors.NewNotFound(sessionResource, sessionID) } +// IsSessionNotFound reports whether err is a not-found error for the session +// resource, rather than a not-found response from a workload or template. +func IsSessionNotFound(err error) bool { + if !apierrors.IsNotFound(err) { + return false + } + var statusErr apierrors.APIStatus + if !errors.As(err, &statusErr) || statusErr.Status().Details == nil { + return false + } + details := statusErr.Status().Details + return details.Group == sessionResource.Group && details.Kind == sessionResource.Resource +} + func workloadResource(kind string) schema.GroupResource { switch kind { case types.CodeInterpreterKind: diff --git a/pkg/api/errors_test.go b/pkg/api/errors_test.go index 92133d359..f136c93ad 100644 --- a/pkg/api/errors_test.go +++ b/pkg/api/errors_test.go @@ -44,6 +44,9 @@ func TestNewSessionNotFoundError(t *testing.T) { assert.Equal(t, sessionResource.Group, status.Details.Group) assert.Equal(t, sessionResource.Resource, status.Details.Kind) assert.Equal(t, sessionID, status.Details.Name) + assert.True(t, IsSessionNotFound(err)) + assert.False(t, IsSessionNotFound(errors.New("not found"))) + assert.False(t, IsSessionNotFound(NewSandboxTemplateNotFoundError("default", "agent", types.AgentRuntimeKind))) } func TestWorkloadResource(t *testing.T) { diff --git a/pkg/common/types/sandbox.go b/pkg/common/types/sandbox.go index c75754e51..4a007f58c 100644 --- a/pkg/common/types/sandbox.go +++ b/pkg/common/types/sandbox.go @@ -55,6 +55,9 @@ type CreateSandboxRequest struct { Kind string `json:"kind"` Name string `json:"name"` Namespace string `json:"namespace"` + // TTL is the client-requested maximum session lifetime in seconds. The + // workload's MaxSessionDuration remains the upper bound. + TTL *int64 `json:"ttl,omitempty"` } type CreateSandboxResponse struct { @@ -79,5 +82,23 @@ func (car *CreateSandboxRequest) Validate() error { if car.Name == "" { return fmt.Errorf("name is required") } + if car.TTL != nil { + const maxDurationSeconds = (1<<63 - 1) / int64(time.Second) + if *car.TTL <= 0 { + return fmt.Errorf("ttl must be greater than zero") + } + if *car.TTL > maxDurationSeconds { + return fmt.Errorf("ttl exceeds the maximum supported duration") + } + } return nil } + +// RequestedTTL returns the requested session lifetime, or zero when the +// client left lifetime selection to the workload configuration. +func (car *CreateSandboxRequest) RequestedTTL() time.Duration { + if car.TTL == nil { + return 0 + } + return time.Duration(*car.TTL) * time.Second +} diff --git a/pkg/common/types/sandbox_test.go b/pkg/common/types/sandbox_test.go index 97294b3d5..6969b694e 100644 --- a/pkg/common/types/sandbox_test.go +++ b/pkg/common/types/sandbox_test.go @@ -18,8 +18,10 @@ package types import ( "testing" + "time" "github.com/stretchr/testify/assert" + "k8s.io/utils/ptr" ) func TestCreateSandboxRequest_Validate(t *testing.T) { @@ -44,9 +46,32 @@ func TestCreateSandboxRequest_Validate(t *testing.T) { Kind: CodeInterpreterKind, Namespace: "default", Name: "test-ci", + TTL: ptr.To[int64](60), }, wantError: false, }, + { + name: "zero ttl", + req: CreateSandboxRequest{ + Kind: CodeInterpreterKind, + Namespace: "default", + Name: "test-ci", + TTL: ptr.To[int64](0), + }, + wantError: true, + errorMsg: "ttl must be greater than zero", + }, + { + name: "ttl overflows time.Duration", + req: CreateSandboxRequest{ + Kind: CodeInterpreterKind, + Namespace: "default", + Name: "test-ci", + TTL: ptr.To[int64](1 << 62), + }, + wantError: true, + errorMsg: "ttl exceeds the maximum supported duration", + }, { name: "invalid kind", req: CreateSandboxRequest{ @@ -171,3 +196,8 @@ func TestCreateSandboxRequest_Validate(t *testing.T) { }) } } + +func TestCreateSandboxRequest_RequestedTTL(t *testing.T) { + assert.Zero(t, (&CreateSandboxRequest{}).RequestedTTL()) + assert.Equal(t, 90*time.Second, (&CreateSandboxRequest{TTL: ptr.To[int64](90)}).RequestedTTL()) +} diff --git a/pkg/router/handlers.go b/pkg/router/handlers.go index 811b87ce3..702591f33 100644 --- a/pkg/router/handlers.go +++ b/pkg/router/handlers.go @@ -31,6 +31,7 @@ import ( apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/klog/v2" + "github.com/volcano-sh/agentcube/pkg/api" "github.com/volcano-sh/agentcube/pkg/common/types" ) @@ -142,7 +143,11 @@ func (s *Server) handleGetSandboxError(c *gin.Context, err error) { if code == http.StatusInternalServerError { message = "internal server error" } - c.JSON(code, gin.H{"error": message}) + response := gin.H{"error": message} + if api.IsSessionNotFound(err) { + response["code"] = api.SessionNotFoundCode + } + c.JSON(code, response) return } diff --git a/pkg/router/handlers_test.go b/pkg/router/handlers_test.go index bd0f76953..e0709e916 100644 --- a/pkg/router/handlers_test.go +++ b/pkg/router/handlers_test.go @@ -18,6 +18,7 @@ package router import ( "context" + "encoding/json" "errors" "fmt" "net" @@ -154,9 +155,10 @@ func TestHandleInvoke_ErrorPaths(t *testing.T) { defer teardownEnv() tests := []struct { - name string - err error - expectedCode int + name string + err error + expectedCode int + expectedErrorCode string }{ { name: "session manager generic error", @@ -164,9 +166,10 @@ func TestHandleInvoke_ErrorPaths(t *testing.T) { expectedCode: http.StatusInternalServerError, }, { - name: "session not found", - err: api.NewSessionNotFoundError("missing-session"), - expectedCode: http.StatusNotFound, + name: "session not found", + err: api.NewSessionNotFoundError("missing-session"), + expectedCode: http.StatusNotFound, + expectedErrorCode: api.SessionNotFoundCode, }, { name: "agent runtime not found", @@ -202,6 +205,13 @@ func TestHandleInvoke_ErrorPaths(t *testing.T) { if w.Code != tt.expectedCode { t.Fatalf("expected status %d, got %d", tt.expectedCode, w.Code) } + var response map[string]string + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if response["code"] != tt.expectedErrorCode { + t.Fatalf("expected error code %q, got %q", tt.expectedErrorCode, response["code"]) + } t.Logf("Response body: %s", w.Body.String()) }) } diff --git a/pkg/workloadmanager/handlers.go b/pkg/workloadmanager/handlers.go index 52f85eb2e..a4ecb2be2 100644 --- a/pkg/workloadmanager/handlers.go +++ b/pkg/workloadmanager/handlers.go @@ -171,9 +171,9 @@ func (s *Server) handleSandboxCreate(c *gin.Context, kind string) { switch sandboxReq.Kind { case types.AgentRuntimeKind: - sandbox, sandboxEntry, err = buildSandboxByAgentRuntime(sandboxReq.Namespace, sandboxReq.Name, ownerID, s.informers) + sandbox, sandboxEntry, err = buildSandboxByAgentRuntime(sandboxReq.Namespace, sandboxReq.Name, ownerID, sandboxReq.RequestedTTL(), s.informers) case types.CodeInterpreterKind: - sandbox, sandboxClaim, sandboxEntry, err = buildSandboxByCodeInterpreter(sandboxReq.Namespace, sandboxReq.Name, ownerID, s.informers) + sandbox, sandboxClaim, sandboxEntry, err = buildSandboxByCodeInterpreter(sandboxReq.Namespace, sandboxReq.Name, ownerID, sandboxReq.RequestedTTL(), s.informers) } if err != nil { diff --git a/pkg/workloadmanager/handlers_test.go b/pkg/workloadmanager/handlers_test.go index 2105967d5..823ccb341 100644 --- a/pkg/workloadmanager/handlers_test.go +++ b/pkg/workloadmanager/handlers_test.go @@ -742,6 +742,7 @@ func TestHandleSandboxCreate(t *testing.T) { expectStatus int expectMessage string expectCreateCalls int + expectTTL time.Duration }{ { name: "invalid json", @@ -829,10 +830,11 @@ func TestHandleSandboxCreate(t *testing.T) { { name: "create sandbox success code interpreter", kind: types.CodeInterpreterKind, - body: `{"name":"workload","namespace":"ns"}`, + body: `{"name":"workload","namespace":"ns","ttl":60}`, createResp: &types.CreateSandboxResponse{SessionID: "sess-1", SandboxID: "id-2", SandboxName: "sandbox-2"}, expectStatus: http.StatusOK, expectCreateCalls: 1, + expectTTL: time.Minute, }, } @@ -853,20 +855,22 @@ func TestHandleSandboxCreate(t *testing.T) { patches := gomonkey.NewPatches() defer patches.Reset() - patches.ApplyFunc(buildSandboxByAgentRuntime, func(_, _, _ string, _ *Informers) (*sandboxv1alpha1.Sandbox, *sandboxEntry, error) { + patches.ApplyFunc(buildSandboxByAgentRuntime, func(_, _, _ string, requestedTTL time.Duration, _ *Informers) (*sandboxv1alpha1.Sandbox, *sandboxEntry, error) { if tc.kind != types.AgentRuntimeKind { return nil, nil, errors.New("unexpected kind") } + require.Equal(t, tc.expectTTL, requestedTTL) if tc.buildErr != nil { return nil, nil, tc.buildErr } return sb, entry, nil }) - patches.ApplyFunc(buildSandboxByCodeInterpreter, func(_, _, _ string, _ *Informers) (*sandboxv1alpha1.Sandbox, *extensionsv1alpha1.SandboxClaim, *sandboxEntry, error) { + patches.ApplyFunc(buildSandboxByCodeInterpreter, func(_, _, _ string, requestedTTL time.Duration, _ *Informers) (*sandboxv1alpha1.Sandbox, *extensionsv1alpha1.SandboxClaim, *sandboxEntry, error) { if tc.kind != types.CodeInterpreterKind { return nil, nil, nil, errors.New("unexpected kind") } + require.Equal(t, tc.expectTTL, requestedTTL) if tc.buildErr != nil { return nil, nil, nil, tc.buildErr } diff --git a/pkg/workloadmanager/workload_builder.go b/pkg/workloadmanager/workload_builder.go index 2a42e4f07..7ebc6315b 100644 --- a/pkg/workloadmanager/workload_builder.go +++ b/pkg/workloadmanager/workload_builder.go @@ -152,6 +152,19 @@ type buildSandboxClaimParams struct { ownerReference *metav1.OwnerReference } +// effectiveSessionTTL applies the workload configuration as the hard upper +// bound for a client-requested session lifetime. +func effectiveSessionTTL(requestedTTL time.Duration, maxSessionDuration *metav1.Duration) time.Duration { + configuredTTL := DefaultSandboxTTL + if maxSessionDuration != nil && maxSessionDuration.Duration > 0 { + configuredTTL = maxSessionDuration.Duration + } + if requestedTTL > 0 && requestedTTL < configuredTTL { + return requestedTTL + } + return configuredTTL +} + // buildSandboxObject builds a Sandbox object from parameters func buildSandboxObject(params *buildSandboxParams) *sandboxv1alpha1.Sandbox { if params.ttl == 0 { @@ -256,7 +269,7 @@ func buildSandboxClaimObject(params *buildSandboxClaimParams) *extensionsv1alpha return sandboxClaim } -func buildSandboxByAgentRuntime(namespace string, name string, ownerID string, ifm *Informers) (*sandboxv1alpha1.Sandbox, *sandboxEntry, error) { +func buildSandboxByAgentRuntime(namespace string, name string, ownerID string, requestedTTL time.Duration, ifm *Informers) (*sandboxv1alpha1.Sandbox, *sandboxEntry, error) { agentRuntimeObj, err := ifm.AgentRuntimeLister.AgentRuntimes(namespace).Get(name) if err != nil { if apierrors.IsNotFound(err) { @@ -280,6 +293,7 @@ func buildSandboxByAgentRuntime(namespace string, name string, ownerID string, i sandboxName: sandboxName, sessionID: sessionID, ownerID: ownerID, + ttl: effectiveSessionTTL(requestedTTL, agentRuntimeObj.Spec.MaxSessionDuration), podSpec: *podSpec, } // Apply labels and annotations from AgentRuntime template @@ -289,9 +303,6 @@ func buildSandboxByAgentRuntime(namespace string, name string, ownerID string, i if agentRuntimeObj.Spec.Template.Annotations != nil { buildParams.podAnnotations = agentRuntimeObj.Spec.Template.Annotations } - if agentRuntimeObj.Spec.MaxSessionDuration != nil { - buildParams.ttl = agentRuntimeObj.Spec.MaxSessionDuration.Duration - } idleTimeout := DefaultSandboxIdleTimeout if agentRuntimeObj.Spec.SessionTimeout != nil { idleTimeout = agentRuntimeObj.Spec.SessionTimeout.Duration @@ -322,7 +333,7 @@ func buildCodeInterpreterEnvVars(templateEnv []corev1.EnvVar, authMode runtimev1 return envVars } -func buildSandboxByCodeInterpreter(namespace string, codeInterpreterName string, ownerID string, informer *Informers) (*sandboxv1alpha1.Sandbox, *extensionsv1alpha1.SandboxClaim, *sandboxEntry, error) { +func buildSandboxByCodeInterpreter(namespace string, codeInterpreterName string, ownerID string, requestedTTL time.Duration, informer *Informers) (*sandboxv1alpha1.Sandbox, *extensionsv1alpha1.SandboxClaim, *sandboxEntry, error) { codeInterpreterObj, err := informer.CodeInterpreterLister.CodeInterpreters(namespace).Get(codeInterpreterName) if err != nil { if apierrors.IsNotFound(err) { @@ -338,6 +349,7 @@ func buildSandboxByCodeInterpreter(namespace string, codeInterpreterName string, sessionID := uuid.New().String() sandboxName := fmt.Sprintf("%s-%s", codeInterpreterName, RandString(8)) + effectiveTTL := effectiveSessionTTL(requestedTTL, codeInterpreterObj.Spec.MaxSessionDuration) idleTimeout := DefaultSandboxIdleTimeout if codeInterpreterObj.Spec.SessionTimeout != nil { @@ -386,10 +398,8 @@ func buildSandboxByCodeInterpreter(namespace string, codeInterpreterName string, }, }, } - if codeInterpreterObj.Spec.MaxSessionDuration != nil { - shutdownTime := metav1.NewTime(time.Now().Add(codeInterpreterObj.Spec.MaxSessionDuration.Duration)) - simpleSandbox.Spec.Lifecycle.ShutdownTime = &shutdownTime - } + shutdownTime := metav1.NewTime(time.Now().Add(effectiveTTL)) + simpleSandbox.Spec.Lifecycle.ShutdownTime = &shutdownTime sandboxEntry.Kind = types.SandboxClaimsKind return simpleSandbox, sandboxClaim, sandboxEntry, nil } @@ -428,10 +438,7 @@ func buildSandboxByCodeInterpreter(namespace string, codeInterpreterName string, podLabels: codeInterpreterObj.Spec.Template.Labels, podAnnotations: codeInterpreterObj.Spec.Template.Annotations, idleTimeout: idleTimeout, - } - - if codeInterpreterObj.Spec.MaxSessionDuration != nil { - buildParams.ttl = codeInterpreterObj.Spec.MaxSessionDuration.Duration + ttl: effectiveTTL, } sandbox := buildSandboxObject(buildParams) return sandbox, nil, sandboxEntry, nil diff --git a/pkg/workloadmanager/workload_builder_test.go b/pkg/workloadmanager/workload_builder_test.go index 035dd5c6b..6160ff565 100644 --- a/pkg/workloadmanager/workload_builder_test.go +++ b/pkg/workloadmanager/workload_builder_test.go @@ -92,6 +92,27 @@ func TestBuildSandboxObject_DoesNotMutateCallerLabels(t *testing.T) { } } +func TestEffectiveSessionTTL(t *testing.T) { + configured := &metav1.Duration{Duration: 4 * time.Hour} + tests := []struct { + name string + requested time.Duration + configured *metav1.Duration + want time.Duration + }{ + {name: "workload default", configured: configured, want: 4 * time.Hour}, + {name: "client ttl below limit", requested: time.Hour, configured: configured, want: time.Hour}, + {name: "client ttl capped", requested: 8 * time.Hour, configured: configured, want: 4 * time.Hour}, + {name: "global default", want: DefaultSandboxTTL}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, effectiveSessionTTL(tt.requested, tt.configured)) + }) + } +} + // TestBuildSandboxObject_NilLabels verifies that a nil podLabels input still // produces a sandbox with the injected session labels. func TestBuildSandboxObject_NilLabels(t *testing.T) { @@ -361,7 +382,7 @@ func TestBuildSandboxByAgentRuntime_NotFound(t *testing.T) { cubeInformerFactory: factory, } - _, _, err := buildSandboxByAgentRuntime(testNamespace, "missing", "", ifm) + _, _, err := buildSandboxByAgentRuntime(testNamespace, "missing", "", 0, ifm) if !errors.Is(err, api.ErrAgentRuntimeNotFound) { t.Fatalf("expected error %v, got %v", api.ErrAgentRuntimeNotFound, err) } @@ -403,7 +424,8 @@ func TestBuildSandboxByAgentRuntime_Success(t *testing.T) { cubeInformerFactory: factory, } - sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", ifm) + startedAt := time.Now() + sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", time.Hour, ifm) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -422,6 +444,7 @@ func TestBuildSandboxByAgentRuntime_Success(t *testing.T) { if sandbox.Spec.Lifecycle.ShutdownTime == nil { t.Error("expected shutdown time to be set") } + assert.WithinDuration(t, startedAt.Add(time.Hour), sandbox.Spec.Lifecycle.ShutdownTime.Time, time.Second) // Validate Entry if entry.Kind != types.SandboxKind { @@ -466,7 +489,7 @@ func TestBuildSandboxByAgentRuntime_DefaultTimeouts(t *testing.T) { cubeInformerFactory: factory, } - sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", ifm) + sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", 0, ifm) assert.NoError(t, err) assert.NotNil(t, sandbox) assert.NotNil(t, entry) @@ -509,7 +532,7 @@ func TestBuildSandboxByAgentRuntime_CustomTimeouts(t *testing.T) { cubeInformerFactory: factory, } - sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", ifm) + sandbox, entry, err := buildSandboxByAgentRuntime(testNamespace, testAgentRuntimeName, "", 0, ifm) assert.NoError(t, err) assert.NotNil(t, sandbox) assert.NotNil(t, entry) @@ -530,7 +553,7 @@ func TestBuildSandboxByCodeInterpreter_NotFound(t *testing.T) { cubeInformerFactory: factory, } - _, _, _, err := buildSandboxByCodeInterpreter(testNamespace, "missing", "", ifm) + _, _, _, err := buildSandboxByCodeInterpreter(testNamespace, "missing", "", 0, ifm) if !errors.Is(err, api.ErrCodeInterpreterNotFound) { t.Fatalf("expected error %v, got %v", api.ErrCodeInterpreterNotFound, err) } @@ -565,7 +588,7 @@ func TestBuildSandboxByCodeInterpreter_PicodAuthFailsWithoutKey(t *testing.T) { cubeInformerFactory: factory, } - _, _, _, err = buildSandboxByCodeInterpreter(testNamespace, "ci-picod-no-key", "", ifm) + _, _, _, err = buildSandboxByCodeInterpreter(testNamespace, "ci-picod-no-key", "", 0, ifm) if !errors.Is(err, api.ErrPublicKeyMissing) { t.Fatalf("expected error %v, got %v", api.ErrPublicKeyMissing, err) } @@ -600,7 +623,8 @@ func TestBuildSandboxByCodeInterpreter_SuccessNoWarmPool(t *testing.T) { cubeInformerFactory: factory, } - sandbox, claim, entry, err := buildSandboxByCodeInterpreter(testNamespace, "ci-no-wp", "", ifm) + startedAt := time.Now() + sandbox, claim, entry, err := buildSandboxByCodeInterpreter(testNamespace, "ci-no-wp", "", 30*time.Minute, ifm) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -614,6 +638,7 @@ func TestBuildSandboxByCodeInterpreter_SuccessNoWarmPool(t *testing.T) { if claim != nil { t.Fatal("expected claim to be nil for non-warm pool path") } + assert.WithinDuration(t, startedAt.Add(30*time.Minute), sandbox.Spec.Lifecycle.ShutdownTime.Time, time.Second) if !strings.HasPrefix(sandbox.Name, "ci-no-wp-") { t.Errorf("expected sandbox name to start with 'ci-no-wp-', got %q", sandbox.Name) @@ -658,7 +683,8 @@ func TestBuildSandboxByCodeInterpreter_SuccessWithWarmPool(t *testing.T) { cubeInformerFactory: factory, } - sandbox, claim, entry, err := buildSandboxByCodeInterpreter(testNamespace, testCodeInterpreterWarmPool, "", ifm) + startedAt := time.Now() + sandbox, claim, entry, err := buildSandboxByCodeInterpreter(testNamespace, testCodeInterpreterWarmPool, "", 30*time.Minute, ifm) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -672,6 +698,7 @@ func TestBuildSandboxByCodeInterpreter_SuccessWithWarmPool(t *testing.T) { if claim == nil { t.Fatal("expected claim not to be nil for warm pool path") } + assert.WithinDuration(t, startedAt.Add(30*time.Minute), sandbox.Spec.Lifecycle.ShutdownTime.Time, time.Second) assertSandboxMetadata(t, sandbox.Labels, sandbox.Name, sandbox.Namespace, testCodeInterpreterWarmPool+"-", "", entry.SessionID) if entry.Kind != types.SandboxClaimsKind { diff --git a/sdk-python/README.md b/sdk-python/README.md index 02f51c7e5..ba1ae87a7 100644 --- a/sdk-python/README.md +++ b/sdk-python/README.md @@ -38,7 +38,7 @@ with CodeInterpreterClient() as client: For long-running applications, you can manually control the lifecycle: ```python -# Create a session with a 1-hour timeout +# Request a session with a 1-hour maximum lifetime client = CodeInterpreterClient(ttl=3600) try: @@ -82,12 +82,17 @@ df.describe().to_csv('/workspace/summary.csv') CodeInterpreterClient( name="custom-template", # CodeInterpreter CRD template name namespace="agentcube", # Kubernetes namespace - ttl=7200, # Session TTL (seconds) + ttl=7200, # Requested maximum lifetime (seconds) session_id="existing-id", # Optional: reuse existing session verbose=True # Enable debug logging ) ``` +If `ttl` is omitted, the CodeInterpreter CRD's `spec.maxSessionDuration` +determines the maximum lifetime. Otherwise, the effective maximum lifetime is +the smaller of `ttl` and `spec.maxSessionDuration`. A session can be reclaimed +earlier when it is idle for `spec.sessionTimeout`. + **Environment Variables**: * `WORKLOAD_MANAGER_URL`: Control Plane URL @@ -112,6 +117,24 @@ client2.run_code("python", "print(open('value.txt').read())") # File persists client2.stop() # Cleanup when done ``` +### Expired Sessions + +When a reused session has been reclaimed, the SDK clears the stale session ID +and raises `SessionNotFoundError`. It does not create a replacement implicitly +because the previous sandbox state has been lost. + +```python +from agentcube import CodeInterpreterClient, SessionNotFoundError + +client = CodeInterpreterClient(session_id=saved_session_id) +try: + client.list_files() +except SessionNotFoundError: + # The old ID is no longer usable. Create a new client only if the + # application can continue without the previous sandbox state. + client = CodeInterpreterClient() +``` + ## Development ```bash diff --git a/sdk-python/agentcube/__init__.py b/sdk-python/agentcube/__init__.py index 496c8ef0d..b87894841 100644 --- a/sdk-python/agentcube/__init__.py +++ b/sdk-python/agentcube/__init__.py @@ -15,6 +15,7 @@ from .code_interpreter import CodeInterpreterClient from .agent_runtime import AgentRuntimeClient from .auth import AuthProvider, TokenAuth, ServiceAccountAuth +from .exceptions import SessionError, SessionNotFoundError __all__ = [ "CodeInterpreterClient", @@ -22,4 +23,6 @@ "AuthProvider", "TokenAuth", "ServiceAccountAuth", + "SessionError", + "SessionNotFoundError", ] diff --git a/sdk-python/agentcube/agent_runtime.py b/sdk-python/agentcube/agent_runtime.py index 9e6cf8b7b..6e746f366 100644 --- a/sdk-python/agentcube/agent_runtime.py +++ b/sdk-python/agentcube/agent_runtime.py @@ -19,6 +19,8 @@ from requests.exceptions import JSONDecodeError from agentcube.auth import AuthProvider from agentcube.clients.agent_runtime_data_plane import AgentRuntimeDataPlaneClient +from agentcube.exceptions import SessionNotFoundError +from agentcube.utils.http import raise_for_session_status from agentcube.utils.log import get_logger @@ -91,7 +93,11 @@ def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None, path: timeout=timeout, path=path, ) - resp.raise_for_status() + try: + raise_for_session_status(resp, self.session_id) + except SessionNotFoundError: + self.session_id = None + raise try: return resp.json() diff --git a/sdk-python/agentcube/clients/code_interpreter_data_plane.py b/sdk-python/agentcube/clients/code_interpreter_data_plane.py index f4c1c0cda..23b75649e 100644 --- a/sdk-python/agentcube/clients/code_interpreter_data_plane.py +++ b/sdk-python/agentcube/clients/code_interpreter_data_plane.py @@ -20,14 +20,14 @@ import os import ast import shlex -from typing import TYPE_CHECKING, Optional, Any, Dict, List, Union +from typing import TYPE_CHECKING, Optional, Any, Callable, Dict, List, Union from urllib.parse import urljoin import requests from agentcube.utils.log import get_logger -from agentcube.utils.http import create_session -from agentcube.exceptions import CommandExecutionError +from agentcube.utils.http import create_session, raise_for_session_status +from agentcube.exceptions import CommandExecutionError, SessionError, SessionNotFoundError if TYPE_CHECKING: from agentcube.auth import AuthProvider @@ -52,6 +52,7 @@ def __init__( pool_connections: int = 10, pool_maxsize: int = 10, auth: Optional["AuthProvider"] = None, + on_session_not_found: Optional[Callable[[], None]] = None, ): """Initialize Data Plane client. @@ -66,12 +67,13 @@ def __init__( pool_connections: Number of connection pools to cache (default: 10). pool_maxsize: Maximum connections per pool (default: 10). """ - self.session_id = session_id + self.session_id: Optional[str] = session_id self.timeout = timeout self.connect_timeout = connect_timeout self.pool_connections = pool_connections self.pool_maxsize = pool_maxsize self._auth = auth + self._on_session_not_found = on_session_not_found self.logger = get_logger(f"{__name__}.CodeInterpreterDataPlaneClient") if base_url: @@ -102,6 +104,9 @@ def _request(self, method: str, endpoint: str, body: Optional[bytes] = None, **k Note: Router handles JWT authentication, so we don't add Authorization header here. """ + if not self.session_id: + raise SessionError("Session is no longer available") + url = urljoin(self.base_url, endpoint) headers = {} @@ -124,13 +129,26 @@ def _request(self, method: str, endpoint: str, body: Optional[bytes] = None, **k self.logger.debug(f"{method} {url}") # Use session for connection pooling - return self.session.request( + response = self.session.request( method=method, url=url, data=body, headers=headers, **kwargs ) + self._raise_for_status(response) + return response + + def _raise_for_status(self, response: requests.Response) -> None: + session_id = self.session_id or "" + try: + raise_for_session_status(response, session_id) + except SessionNotFoundError: + self.session_id = None + self.session.headers.pop("x-agentcube-session-id", None) + if self._on_session_not_found: + self._on_session_not_found() + raise def execute_command_result( self, command: Union[str, List[str]], timeout: Optional[float] = None @@ -153,8 +171,6 @@ def execute_command_result( read_timeout = timeout_value + 2.0 if isinstance(timeout_value, (int, float)) else timeout_value resp = self._request("POST", "api/execute", body=body, timeout=read_timeout) - resp.raise_for_status() - result = resp.json() return { "stdout": result.get("stdout") or "", @@ -228,13 +244,14 @@ def write_file(self, content: str, remote_path: str) -> None: } body = json.dumps(payload).encode('utf-8') - resp = self._request("POST", "api/files", body=body) - resp.raise_for_status() + self._request("POST", "api/files", body=body) def upload_file(self, local_path: str, remote_path: str) -> None: """Upload a local file using multipart/form-data.""" if not os.path.exists(local_path): raise FileNotFoundError(f"Local file not found: {local_path}") + if not self.session_id: + raise SessionError("Session is no longer available") with open(local_path, 'rb') as f: files = {'file': f} @@ -249,14 +266,12 @@ def upload_file(self, local_path: str, remote_path: str) -> None: self.logger.debug(f"Uploading file {local_path} to {remote_path}") resp = self.session.post(url, files=files, data=data, headers=headers, timeout=self.timeout) - resp.raise_for_status() + self._raise_for_status(resp) def download_file(self, remote_path: str, local_path: str) -> None: """Download a file.""" clean_path = remote_path.lstrip("/") resp = self._request("GET", f"api/files/{clean_path}", stream=True) - resp.raise_for_status() - if os.path.dirname(local_path): os.makedirs(os.path.dirname(local_path), exist_ok=True) with open(local_path, 'wb') as f: @@ -266,7 +281,6 @@ def download_file(self, remote_path: str, local_path: str) -> None: def list_files(self, path: str = ".") -> Any: """List files in a directory.""" resp = self._request("GET", "api/files", params={"path": path}) - resp.raise_for_status() return resp.json().get("files", []) def close(self): diff --git a/sdk-python/agentcube/clients/control_plane.py b/sdk-python/agentcube/clients/control_plane.py index 964dffd8b..fc26aa8c5 100644 --- a/sdk-python/agentcube/clients/control_plane.py +++ b/sdk-python/agentcube/clients/control_plane.py @@ -103,7 +103,7 @@ def create_session( name: str = "my-interpreter", namespace: str = "default", metadata: Optional[Dict[str, Any]] = None, - ttl: int = 3600, + ttl: Optional[int] = None, ) -> str: """Create a new Code Interpreter session. @@ -111,17 +111,19 @@ def create_session( name: Name of the CodeInterpreter template (CRD name). namespace: Kubernetes namespace. metadata: Optional metadata. - ttl: Time to live (seconds). + ttl: Optional requested maximum lifetime in seconds. When omitted, + the workload configuration determines the lifetime. Returns: session_id (str): The ID of the created session. """ - payload = { + payload: Dict[str, Any] = { "name": name, "namespace": namespace, - "ttl": ttl, "metadata": metadata or {} } + if ttl is not None: + payload["ttl"] = ttl url = f"{self.base_url}/v1/code-interpreter" self.logger.debug(f"Creating session at {url} with payload: {payload}") diff --git a/sdk-python/agentcube/code_interpreter.py b/sdk-python/agentcube/code_interpreter.py index ca560627c..07e925068 100644 --- a/sdk-python/agentcube/code_interpreter.py +++ b/sdk-python/agentcube/code_interpreter.py @@ -20,6 +20,7 @@ from agentcube.clients.control_plane import ControlPlaneClient from agentcube.clients.code_interpreter_data_plane import CodeInterpreterDataPlaneClient +from agentcube.exceptions import SessionError from agentcube.utils.log import get_logger if TYPE_CHECKING: @@ -57,7 +58,7 @@ def __init__( self, name: str = "my-interpreter", namespace: str = "default", - ttl: int = 3600, + ttl: Optional[int] = None, workload_manager_url: Optional[str] = None, router_url: Optional[str] = None, auth_token: Optional[str] = None, @@ -74,7 +75,8 @@ def __init__( Args: name: Name of the CodeInterpreter template (CRD name). namespace: Kubernetes namespace. - ttl: Time to live (seconds) for new sessions. + ttl: Optional requested maximum lifetime in seconds for new sessions. + When omitted, the workload configuration determines the lifetime. workload_manager_url: URL of WorkloadManager (Control Plane). router_url: URL of Router (Data Plane). auth_token: Auth token for Kubernetes/WorkloadManager. @@ -149,10 +151,20 @@ def _init_data_plane(self): namespace=self.namespace, session_id=self.session_id, auth=self._auth, + on_session_not_found=self._invalidate_session, ) if self.verbose: self.dp_client.logger.setLevel(logging.DEBUG) + def _invalidate_session(self): + """Forget a session that the Router reports as missing.""" + self.session_id = None + + def _require_data_plane(self) -> CodeInterpreterDataPlaneClient: + if not self.session_id or not self.dp_client: + raise SessionError("Session is no longer available; create a new client to continue") + return self.dp_client + def __enter__(self): return self @@ -190,13 +202,13 @@ def execute_command(self, command: str, timeout: Optional[float] = None) -> str: Returns: str: The output of the command. """ - return self.dp_client.execute_command(command, timeout) + return self._require_data_plane().execute_command(command, timeout) def execute_command_result( self, command: str, timeout: Optional[float] = None ) -> dict[str, Any]: """Run a shell command and return ``stdout``, ``stderr``, and ``exit_code`` (no raise on failure).""" - return self.dp_client.execute_command_result(command, timeout) + return self._require_data_plane().execute_command_result(command, timeout) def run_code(self, language: str, code: str, timeout: Optional[float] = None) -> str: """ @@ -214,7 +226,7 @@ def run_code(self, language: str, code: str, timeout: Optional[float] = None) -> Returns: The standard output (stdout) generated by the code execution. """ - return self.dp_client.run_code(language, code, timeout) + return self._require_data_plane().run_code(language, code, timeout) def write_file(self, content: str, remote_path: str): """ @@ -225,7 +237,7 @@ def write_file(self, content: str, remote_path: str): remote_path: The destination path of the file in the remote environment. This path is relative to the session's working directory. """ - self.dp_client.write_file(content, remote_path) + self._require_data_plane().write_file(content, remote_path) def upload_file(self, local_path: str, remote_path: str): """ @@ -236,7 +248,7 @@ def upload_file(self, local_path: str, remote_path: str): remote_path: The destination path of the file in the remote environment. This path is relative to the session's working directory. """ - self.dp_client.upload_file(local_path, remote_path) + self._require_data_plane().upload_file(local_path, remote_path) def download_file(self, remote_path: str, local_path: str): """ @@ -247,7 +259,7 @@ def download_file(self, remote_path: str, local_path: str): This path is relative to the session's working directory. local_path: The destination path on the local filesystem to save the file. """ - self.dp_client.download_file(remote_path, local_path) + self._require_data_plane().download_file(remote_path, local_path) def list_files(self, path: str = "."): """ @@ -259,4 +271,4 @@ def list_files(self, path: str = "."): Returns: A list of file/directory information dicts. """ - return self.dp_client.list_files(path) + return self._require_data_plane().list_files(path) diff --git a/sdk-python/agentcube/exceptions.py b/sdk-python/agentcube/exceptions.py index 0240584f4..cba27508c 100644 --- a/sdk-python/agentcube/exceptions.py +++ b/sdk-python/agentcube/exceptions.py @@ -28,6 +28,13 @@ class SessionError(AgentCubeError): """Raised when session creation or management fails""" pass +class SessionNotFoundError(SessionError): + """Raised when the server no longer has the requested session.""" + def __init__(self, session_id, message=None, response=None): + self.session_id = session_id + self.response = response + super().__init__(message or f"Session {session_id!r} was not found") + class DataPlaneError(AgentCubeError): """Raised when Data Plane operations fail""" pass diff --git a/sdk-python/agentcube/utils/http.py b/sdk-python/agentcube/utils/http.py index fefed028f..978464ae6 100644 --- a/sdk-python/agentcube/utils/http.py +++ b/sdk-python/agentcube/utils/http.py @@ -17,6 +17,11 @@ import requests from requests.adapters import HTTPAdapter +from agentcube.exceptions import SessionNotFoundError + + +SESSION_NOT_FOUND_CODE = "SESSION_NOT_FOUND" + def create_session( pool_connections: int = 10, @@ -42,3 +47,19 @@ def create_session( session.mount("https://", adapter) return session + + +def raise_for_session_status(response: requests.Response, session_id: str) -> None: + """Raise a typed error for a missing session, or the normal HTTP error.""" + if response.status_code == 404: + try: + payload = response.json() + except (TypeError, ValueError): + payload = {} + if isinstance(payload, dict) and payload.get("code") == SESSION_NOT_FOUND_CODE: + raise SessionNotFoundError( + session_id=session_id, + message=payload.get("error"), + response=response, + ) + response.raise_for_status() diff --git a/sdk-python/tests/test_agent_runtime.py b/sdk-python/tests/test_agent_runtime.py index 849c27780..c63b96b39 100644 --- a/sdk-python/tests/test_agent_runtime.py +++ b/sdk-python/tests/test_agent_runtime.py @@ -20,6 +20,7 @@ os.environ.setdefault("ROUTER_URL", "http://mock-router:8080") from agentcube.agent_runtime import AgentRuntimeClient +from agentcube.exceptions import SessionNotFoundError class TestAgentRuntimeClientSessionBootstrap(unittest.TestCase): @@ -89,6 +90,45 @@ def test_invoke_falls_back_to_text_when_non_json(self, mock_dp_class): self.assertEqual(out, "plain") + @patch("agentcube.agent_runtime.AgentRuntimeDataPlaneClient") + def test_session_not_found_invalidates_session(self, mock_dp_class): + mock_dp = Mock() + mock_dp.bootstrap_session_id.return_value = "expired-session" + + resp = Mock() + resp.status_code = 404 + resp.json.return_value = { + "code": "SESSION_NOT_FOUND", + "error": "session expired-session was not found", + } + mock_dp.invoke.return_value = resp + mock_dp_class.return_value = mock_dp + + client = AgentRuntimeClient(agent_name="agent-a", router_url="http://t:1") + with self.assertRaises(SessionNotFoundError) as ctx: + client.invoke({"input": "hi"}) + + self.assertEqual(ctx.exception.session_id, "expired-session") + self.assertIsNone(client.session_id) + + @patch("agentcube.agent_runtime.AgentRuntimeDataPlaneClient") + def test_application_404_does_not_invalidate_session(self, mock_dp_class): + mock_dp = Mock() + mock_dp.bootstrap_session_id.return_value = "active-session" + + resp = Mock() + resp.status_code = 404 + resp.json.return_value = {"error": "application route not found"} + resp.raise_for_status.side_effect = requests.exceptions.HTTPError(response=resp) + mock_dp.invoke.return_value = resp + mock_dp_class.return_value = mock_dp + + client = AgentRuntimeClient(agent_name="agent-a", router_url="http://t:1") + with self.assertRaises(requests.exceptions.HTTPError): + client.invoke({"input": "hi"}) + + self.assertEqual(client.session_id, "active-session") + class TestAgentRuntimeDataPlaneClient(unittest.TestCase): @patch("agentcube.clients.agent_runtime_data_plane.create_session") diff --git a/sdk-python/tests/test_code_interpreter.py b/sdk-python/tests/test_code_interpreter.py index 0dfc5d3ba..6cf732fda 100644 --- a/sdk-python/tests/test_code_interpreter.py +++ b/sdk-python/tests/test_code_interpreter.py @@ -26,10 +26,14 @@ import unittest from unittest.mock import Mock, patch +import requests + # Set required env var before import os.environ.setdefault("ROUTER_URL", "http://mock-router:8080") from agentcube.code_interpreter import CodeInterpreterClient +from agentcube.clients.control_plane import ControlPlaneClient +from agentcube.exceptions import SessionError, SessionNotFoundError class TestCodeInterpreterClientInit(unittest.TestCase): @@ -47,9 +51,28 @@ def test_init_creates_session(self, mock_cp_class, mock_dp_class): # Session should be created self.assertEqual(client.session_id, "new-session-123") - mock_cp.create_session.assert_called_once() + mock_cp.create_session.assert_called_once_with( + name="my-interpreter", + namespace="default", + ttl=None, + ) mock_dp_class.assert_called_once() + @patch('agentcube.code_interpreter.CodeInterpreterDataPlaneClient') + @patch('agentcube.code_interpreter.ControlPlaneClient') + def test_init_forwards_explicit_ttl(self, mock_cp_class, mock_dp_class): + mock_cp = Mock() + mock_cp.create_session.return_value = "new-session-123" + mock_cp_class.return_value = mock_cp + + CodeInterpreterClient(router_url="http://test:8080", ttl=600) + + mock_cp.create_session.assert_called_once_with( + name="my-interpreter", + namespace="default", + ttl=600, + ) + @patch('agentcube.code_interpreter.CodeInterpreterDataPlaneClient') @patch('agentcube.code_interpreter.ControlPlaneClient') def test_init_with_session_id_reuses_session(self, mock_cp_class, mock_dp_class): @@ -68,6 +91,38 @@ def test_init_with_session_id_reuses_session(self, mock_cp_class, mock_dp_class) mock_dp_class.assert_called_once() +class TestControlPlaneClientTTL(unittest.TestCase): + @patch('agentcube.clients.control_plane.create_session') + def test_omits_ttl_by_default(self, mock_create_session): + session = Mock() + response = Mock() + response.json.return_value = {"sessionId": "new-session-123"} + session.post.return_value = response + session.headers = {} + mock_create_session.return_value = session + + client = ControlPlaneClient(workload_manager_url="http://test:8080") + client.create_session() + + payload = session.post.call_args.kwargs["json"] + self.assertNotIn("ttl", payload) + + @patch('agentcube.clients.control_plane.create_session') + def test_includes_explicit_ttl(self, mock_create_session): + session = Mock() + response = Mock() + response.json.return_value = {"sessionId": "new-session-123"} + session.post.return_value = response + session.headers = {} + mock_create_session.return_value = session + + client = ControlPlaneClient(workload_manager_url="http://test:8080") + client.create_session(ttl=600) + + payload = session.post.call_args.kwargs["json"] + self.assertEqual(payload["ttl"], 600) + + class TestSessionIdProperty(unittest.TestCase): """Test session_id property.""" @@ -108,6 +163,25 @@ def test_reuse_session_no_new_creation(self, mock_cp_class, mock_dp_class): call_kwargs = mock_dp_class.call_args[1] self.assertEqual(call_kwargs['session_id'], "reused-session-789") + @patch('agentcube.code_interpreter.CodeInterpreterDataPlaneClient') + @patch('agentcube.code_interpreter.ControlPlaneClient') + def test_missing_session_is_invalidated(self, mock_cp_class, mock_dp_class): + mock_cp_class.return_value = Mock() + mock_dp = Mock() + mock_dp_class.return_value = mock_dp + + client = CodeInterpreterClient( + router_url="http://test:8080", + session_id="expired-session", + ) + callback = mock_dp_class.call_args.kwargs["on_session_not_found"] + callback() + + self.assertIsNone(client.session_id) + with self.assertRaises(SessionError): + client.list_files() + mock_dp.list_files.assert_not_called() + class TestContextManager(unittest.TestCase): """Test context manager behavior.""" @@ -155,5 +229,63 @@ def test_cleanup_on_dp_init_failure(self, mock_cp_class, mock_dp_class): mock_cp.delete_session.assert_called_once_with("leaked-session-999") +class TestCodeInterpreterDataPlaneSessionErrors(unittest.TestCase): + @patch('agentcube.clients.code_interpreter_data_plane.create_session') + def test_session_not_found_raises_typed_error_and_invalidates(self, mock_create_session): + session = Mock() + response = requests.Response() + response.status_code = 404 + response._content = b'{"code":"SESSION_NOT_FOUND","error":"session expired"}' + session.request.return_value = response + session.headers = requests.structures.CaseInsensitiveDict({ + "x-agentcube-session-id": "expired-session", + }) + mock_create_session.return_value = session + callback = Mock() + + from agentcube.clients.code_interpreter_data_plane import CodeInterpreterDataPlaneClient + + client = CodeInterpreterDataPlaneClient( + session_id="expired-session", + base_url="http://router/invocations/", + on_session_not_found=callback, + ) + + with self.assertRaises(SessionNotFoundError) as ctx: + client.list_files() + + self.assertEqual(ctx.exception.session_id, "expired-session") + self.assertIsNone(client.session_id) + self.assertNotIn("x-agentcube-session-id", session.headers) + callback.assert_called_once_with() + with self.assertRaises(SessionError): + client.list_files() + session.request.assert_called_once() + + @patch('agentcube.clients.code_interpreter_data_plane.create_session') + def test_application_404_remains_http_error(self, mock_create_session): + session = Mock() + response = requests.Response() + response.status_code = 404 + response._content = b'{"error":"application route not found"}' + session.request.return_value = response + session.headers = requests.structures.CaseInsensitiveDict({ + "x-agentcube-session-id": "active-session", + }) + mock_create_session.return_value = session + + from agentcube.clients.code_interpreter_data_plane import CodeInterpreterDataPlaneClient + + client = CodeInterpreterDataPlaneClient( + session_id="active-session", + base_url="http://router/invocations/", + ) + + with self.assertRaises(requests.exceptions.HTTPError): + client.list_files() + + self.assertEqual(client.session_id, "active-session") + + if __name__ == "__main__": unittest.main()