Skip to content

Commit 87fcbe4

Browse files
taltultctaltultc
committed
feat(tool): change customer
1 parent fe63bae commit 87fcbe4

File tree

5 files changed

+312
-4
lines changed

5 files changed

+312
-4
lines changed

‎doit-mcp-server/src/app.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ app.get("/authorize", async (c) => {
6969

7070
// Reusable approve handler function
7171
async function handleApprove(c: any) {
72-
const { action, oauthReqInfo, apiKey, customerContext } =
72+
const { action, oauthReqInfo, apiKey, customerContext, isDoitUser } =
7373
await parseApproveFormBody(await c.req.parseBody());
7474

7575
if (!oauthReqInfo) {
@@ -88,6 +88,7 @@ async function handleApprove(c: any) {
8888
props: {
8989
apiKey,
9090
customerContext,
91+
isDoitUser,
9192
},
9293
});
9394

‎doit-mcp-server/src/index.ts

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ import {
4848
listInvoicesTool,
4949
getInvoiceTool,
5050
} from "../../src/tools/invoices.js";
51+
import {
52+
ChangeCustomerArgumentsSchema,
53+
changeCustomerTool,
54+
} from "../../src/tools/changeCustomer.js";
5155

5256
import OAuthProvider from "@cloudflare/workers-oauth-provider";
5357
import { executeToolHandler } from "../../src/utils/toolsHandler.js";
@@ -76,10 +80,37 @@ export class DoitMCPAgent extends McpAgent {
7680
return this.props.apiKey as string;
7781
}
7882

83+
// Persist props to Durable Object storage
84+
private async saveProps(): Promise<void> {
85+
if (this.ctx?.storage) {
86+
await this.ctx.storage.put("persistedProps", {
87+
customerContext: this.props.customerContext,
88+
lastUpdated: Date.now(),
89+
});
90+
}
91+
}
92+
93+
// Load props from Durable Object storage
94+
private async loadPersistedProps(): Promise<void> {
95+
if (this.ctx?.storage) {
96+
const persistedProps = await this.ctx.storage.get<{
97+
apiKey: string;
98+
customerContext: string;
99+
lastUpdated: number;
100+
}>("persistedProps");
101+
if (persistedProps) {
102+
console.log("Loading persisted props:", persistedProps);
103+
// Update props with persisted values
104+
this.props.customerContext = persistedProps.customerContext;
105+
}
106+
}
107+
}
108+
79109
// Generic callback factory for tools
80110
private createToolCallback(toolName: string) {
81111
return async (args: any) => {
82112
const token = this.getToken();
113+
83114
const argsWithCustomerContext = {
84115
...args,
85116
customerContext: this.props.customerContext,
@@ -93,6 +124,32 @@ export class DoitMCPAgent extends McpAgent {
93124
};
94125
}
95126

127+
// Special callback for changeCustomer tool
128+
private createChangeCustomerCallback() {
129+
return async (args: any) => {
130+
const token = this.getToken();
131+
const { handleChangeCustomerRequest } = await import(
132+
"../../src/tools/changeCustomer.js"
133+
);
134+
135+
// Create update function to modify the customer context
136+
const updateCustomerContext = async (newContext: string) => {
137+
this.props.customerContext = newContext;
138+
139+
// Persist the updated props
140+
await this.saveProps();
141+
};
142+
143+
const response = await handleChangeCustomerRequest(
144+
args,
145+
token,
146+
updateCustomerContext
147+
);
148+
149+
return convertToMcpResponse(response);
150+
};
151+
}
152+
96153
// Generic tool registration helper
97154
private registerTool(tool: any, schema: any) {
98155
(this.server.tool as any)(
@@ -106,6 +163,16 @@ export class DoitMCPAgent extends McpAgent {
106163
async init() {
107164
console.log("Initializing Doit MCP Agent", this.props.customerContext);
108165

166+
// Load persisted props first
167+
await this.loadPersistedProps();
168+
169+
console.log("After loading persisted props:", this.props.customerContext);
170+
171+
// Save current props if they exist (for initial OAuth setup)
172+
if (this.props.apiKey && this.props.customerContext) {
173+
await this.saveProps();
174+
}
175+
109176
// Register prompts
110177
prompts.forEach((prompt) => {
111178
this.server.prompt(prompt.name, prompt.description, async () => ({
@@ -148,6 +215,16 @@ export class DoitMCPAgent extends McpAgent {
148215
// Invoices tools
149216
this.registerTool(listInvoicesTool, ListInvoicesArgumentsSchema);
150217
this.registerTool(getInvoiceTool, GetInvoiceArgumentsSchema);
218+
219+
// Change Customer tool (requires special handling)
220+
if (this.props.isDoitUser === "true") {
221+
(this.server.tool as any)(
222+
changeCustomerTool.name,
223+
changeCustomerTool.description,
224+
zodSchemaToMcpTool(ChangeCustomerArgumentsSchema),
225+
this.createChangeCustomerCallback()
226+
);
227+
}
151228
}
152229
}
153230

‎doit-mcp-server/src/utils.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -297,15 +297,15 @@ export const parseApproveFormBody = async (body: {
297297
const action = body.action as string;
298298
const apiKey = body.apiKey as string;
299299
const customerContext = body.customerContext as string;
300-
300+
const isDoitUser = body.isDoitUser as string;
301301
let oauthReqInfo: AuthRequest | null = null;
302302
try {
303303
oauthReqInfo = JSON.parse(body.oauthReqInfo as string) as AuthRequest;
304304
} catch (e) {
305305
oauthReqInfo = null;
306306
}
307307

308-
return { action, oauthReqInfo, apiKey, customerContext };
308+
return { action, oauthReqInfo, apiKey, customerContext, isDoitUser };
309309
};
310310

311311
export const renderCustomerContextScreen = async (
@@ -335,7 +335,7 @@ export const renderCustomerContextScreen = async (
335335
value="${JSON.stringify(oauthReqInfo)}"
336336
/>
337337
<input type="hidden" name="apiKey" value="${apiKey}" />
338-
<input type="hidden" name="action" value="approve" />
338+
<input type="hidden" name="isDoitUser" value="true" />
339339
<input
340340
type="text"
341341
name="customerContext"

‎src/tools/__tests__/changeCustomer.test.ts

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest";
2+
import { z } from "zod";
3+
import {
4+
ChangeCustomerArgumentsSchema,
5+
handleChangeCustomerRequest,
6+
} from "../changeCustomer.js";
7+
import {
8+
createErrorResponse,
9+
createSuccessResponse,
10+
formatZodError,
11+
handleGeneralError,
12+
} from "../../utils/util.js";
13+
14+
// Mock the utility functions
15+
vi.mock("../../utils/util.js", () => ({
16+
createErrorResponse: vi.fn((message) => ({
17+
content: [{ type: "text", text: message }],
18+
})),
19+
createSuccessResponse: vi.fn((text) => ({
20+
content: [{ type: "text", text }],
21+
})),
22+
formatZodError: vi.fn((error) => `Validation error: ${error.message}`),
23+
handleGeneralError: vi.fn((error, context) => ({
24+
content: [{ type: "text", text: `Error in ${context}: ${error.message}` }],
25+
})),
26+
}));
27+
28+
describe("changeCustomer", () => {
29+
beforeEach(() => {
30+
vi.clearAllMocks();
31+
});
32+
33+
describe("ChangeCustomerArgumentsSchema", () => {
34+
it("should validate valid arguments", () => {
35+
const validArgs = {
36+
customerContext: "new-customer-123",
37+
};
38+
39+
const result = ChangeCustomerArgumentsSchema.parse(validArgs);
40+
expect(result).toEqual(validArgs);
41+
});
42+
43+
it("should reject invalid arguments", () => {
44+
const invalidArgs = {
45+
// missing customerContext
46+
};
47+
48+
expect(() => ChangeCustomerArgumentsSchema.parse(invalidArgs)).toThrow();
49+
});
50+
51+
it("should reject non-string customerContext", () => {
52+
const invalidArgs = {
53+
customerContext: 123,
54+
};
55+
56+
expect(() => ChangeCustomerArgumentsSchema.parse(invalidArgs)).toThrow();
57+
});
58+
});
59+
60+
describe("handleChangeCustomerRequest", () => {
61+
it("should successfully change customer context", async () => {
62+
const args = {
63+
customerContext: "old-customer",
64+
};
65+
const newContext = "new-customer-123";
66+
const token = "mock-token";
67+
const updateCallback = vi.fn();
68+
69+
const result = await handleChangeCustomerRequest(
70+
{ ...args, customerContext: newContext },
71+
token,
72+
updateCallback
73+
);
74+
75+
expect(updateCallback).toHaveBeenCalledWith(newContext);
76+
expect(createSuccessResponse).toHaveBeenCalledWith(
77+
"Customer context successfully changed from 'old-customer' to 'new-customer-123'"
78+
);
79+
});
80+
81+
it("should handle missing previous context", async () => {
82+
const args = {};
83+
const newContext = "new-customer-123";
84+
const token = "mock-token";
85+
const updateCallback = vi.fn();
86+
87+
const result = await handleChangeCustomerRequest(
88+
{ ...args, customerContext: newContext },
89+
token,
90+
updateCallback
91+
);
92+
93+
expect(updateCallback).toHaveBeenCalledWith(newContext);
94+
expect(createSuccessResponse).toHaveBeenCalledWith(
95+
"Customer context successfully changed to 'new-customer-123'"
96+
);
97+
});
98+
99+
it("should work without update callback", async () => {
100+
const args = {
101+
customerContext: "old-customer",
102+
};
103+
const newContext = "new-customer-123";
104+
const token = "mock-token";
105+
106+
const result = await handleChangeCustomerRequest(
107+
{ ...args, customerContext: newContext },
108+
token
109+
);
110+
111+
expect(createSuccessResponse).toHaveBeenCalledWith(
112+
"Customer context successfully changed from 'old-customer' to 'new-customer-123'"
113+
);
114+
});
115+
116+
it("should handle validation errors", async () => {
117+
const invalidArgs = {
118+
customerContext: 123, // invalid type
119+
};
120+
const token = "mock-token";
121+
122+
const result = await handleChangeCustomerRequest(invalidArgs, token);
123+
124+
expect(createErrorResponse).toHaveBeenCalled();
125+
expect(formatZodError).toHaveBeenCalled();
126+
});
127+
128+
it("should handle general errors", async () => {
129+
// Mock ChangeCustomerArgumentsSchema.parse to throw a non-Zod error
130+
const originalParse = ChangeCustomerArgumentsSchema.parse;
131+
vi.spyOn(ChangeCustomerArgumentsSchema, "parse").mockImplementation(
132+
() => {
133+
throw new Error("Unexpected error");
134+
}
135+
);
136+
137+
const args = {
138+
customerContext: "new-customer-123",
139+
};
140+
const token = "mock-token";
141+
142+
const result = await handleChangeCustomerRequest(args, token);
143+
144+
expect(handleGeneralError).toHaveBeenCalledWith(
145+
expect.any(Error),
146+
"handling change customer request"
147+
);
148+
149+
// Restore original implementation
150+
ChangeCustomerArgumentsSchema.parse = originalParse;
151+
});
152+
});
153+
});

‎src/tools/changeCustomer.ts

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
import { z } from "zod";
2+
import {
3+
createErrorResponse,
4+
createSuccessResponse,
5+
formatZodError,
6+
handleGeneralError,
7+
} from "../utils/util.js";
8+
9+
// Schema definition
10+
export const ChangeCustomerArgumentsSchema = z.object({
11+
customerContext: z.string().describe("The new customer context to set"),
12+
});
13+
14+
// Interfaces
15+
export interface ChangeCustomerResponse {
16+
success: boolean;
17+
previousContext?: string;
18+
newContext: string;
19+
message: string;
20+
}
21+
22+
// Tool metadata
23+
export const changeCustomerTool = {
24+
name: "change_customer",
25+
description:
26+
"Changes the current customer context for subsequent API calls. This allows switching between different customer accounts or contexts, Example: EE8CtpzYiKp0dVAESVrB",
27+
inputSchema: {
28+
type: "object",
29+
properties: {
30+
customerContext: {
31+
type: "string",
32+
description: "The new customer context to set",
33+
},
34+
},
35+
required: ["customerContext"],
36+
},
37+
};
38+
39+
// Handle change customer request
40+
export async function handleChangeCustomerRequest(
41+
args: any,
42+
token: string,
43+
updateCustomerContext?: (newContext: string) => Promise<void> | void
44+
) {
45+
try {
46+
// Validate arguments
47+
const validatedArgs = ChangeCustomerArgumentsSchema.parse(args);
48+
const { customerContext: newContext } = validatedArgs;
49+
if (!newContext) {
50+
return createErrorResponse("Customer context is required");
51+
}
52+
53+
const previousContext = args.customerContext;
54+
55+
// Update the customer context if callback is provided
56+
if (updateCustomerContext) {
57+
await updateCustomerContext(newContext);
58+
}
59+
60+
// Create response
61+
const response: ChangeCustomerResponse = {
62+
success: true,
63+
previousContext,
64+
newContext,
65+
message: `Customer context successfully changed${
66+
previousContext ? ` from '${previousContext}'` : ""
67+
} to '${newContext}'`,
68+
};
69+
70+
return createSuccessResponse(response.message);
71+
} catch (error) {
72+
if (error instanceof z.ZodError) {
73+
return createErrorResponse(formatZodError(error));
74+
}
75+
return handleGeneralError(error, "handling change customer request");
76+
}
77+
}

0 commit comments

Comments
 (0)